Fix PyPI agent apply patching only one installed copy (#529, #501) (#538)

* Start fix for #529, #501

Assisted-by: Claude Code:claude-opus-5-5

* Test PyPI apply patches every installed copy

Regression tests for #529 (Pipenv WORKON_HOME venv beside ./.venv)
and #501 (user site beside a system dist-packages in global scope):
agent apply must patch both copies, vex must then attest, and
rollback must restore both.

Assisted-by: Claude Code:claude-opus-5-5

* Patch every PyPI copy in agent apply

The Python crawler can find one release in several site-packages
dirs, but apply patched only the first. The copy the interpreter
actually imports could stay vulnerable while vex attested it.
Apply now patches every distinct copy, like gem and npm; paths
that alias one directory through a symlink count once.

Fixes #529, #501

Assisted-by: Claude Code:claude-opus-5-5

* Assert per-copy applied events in multi-copy test

The apply JSON envelope reports per-copy outcomes as events, not
results; count the applied events for the patched release.

Assisted-by: Claude Code:claude-opus-5-5

* Prefetch blobs for each copy's PyPI variant

Two environments can hold different wheels of one release. Apply
gates each variant per copy, but the blob prefetch gated it against
the first copy only, so a locally modified file in the second copy
never got its full patched blob and its apply then failed. The
prefetch now probes each variant on the copies it is applied to.

Assisted-by: Claude Code:claude-opus-5-5

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: mikolalysenko

CHANGELOG.md

@@ -102,6 +102,11 @@ limits, and required install commands.
 
 ### Fixed
 
+- Agent-mode PyPI `apply` patches every installed copy of a release, not just
+  the first one found. A Pipenv project with both a WORKON_HOME venv and a
+  `./.venv`, or a global install with the same release in the user site and a
+  system dir, no longer keeps the copy Python imports unpatched while `vex`
+  attests it (#529, #501).
 - Gem hosted and vendored modes wire only the manifest Bundler loads. A `gems.rb`
   twin or a `BUNDLE_GEMFILE` setting (environment or `.bundle/config`) no longer
   leads to an edit of an ignored `Gemfile` that reports success and attests an

crates/socket-patch-cli/CLI_CONTRACT.md

@@ -502,7 +502,11 @@ the model is **not uniform** today:
   **coexisting physical copies of one `gem@version`** (bundler-2's scoped store beside bundler-1's
   flat store), `apply`/`rollback` patch/restore **every copy** — one summary event per copy,
   mirroring npm's multi-copy fan-out — while single-representative consumers (`get`, `vendor`,
-  `vex`) use the highest-precedence copy.
+  `vex`) use the highest-precedence copy. PyPI follows the same rule: when the crawler resolves
+  one release in several site-packages dirs (a Pipenv WORKON_HOME venv beside an auto-detected
+  `./.venv`, or the user site beside a system dir in global scope), agent `apply` patches **every**
+  copy, one summary event per copy, because any of them may be the one the interpreter imports.
+  Paths that resolve to the same directory (a symlinked site-packages) count as one copy.
 
   *Copy classes (additive to the multi-copy vocabulary):* a copy under a **bundle-path store**
   (config/env/default root) is PRIMARY — a variant mismatch or write failure there fails the run,

crates/socket-patch-cli/src/commands/apply.rs

@@ -207,10 +207,11 @@ fn format_mismatch_fetch_result(downloaded: usize, needed: usize) -> String {
 /// duplicates of one `name@version`, the apply loop patches each of them,
 /// and copies drift independently — a pristine (or already-patched) root
 /// copy says nothing about a locally-modified nested duplicate, whose
-/// mismatched files still need their afterHash blobs. The variant gate,
-/// by contrast, mirrors the apply loop's representative check against the
-/// FIRST copy (release-variant ecosystems install one directory per
-/// `package@version`).
+/// mismatched files still need their afterHash blobs. The variant gate
+/// mirrors the apply loop's representative check PER COPY for gem and
+/// PyPI (which patch every copy, and two envs can hold different wheels
+/// of one release), and against the FIRST copy for Maven: a variant's
+/// files are probed only on the copies it is attempted on.
 ///
 /// Only a mismatched file whose afterHash blob is NOT staged can queue a
 /// fetch, so the probe first decides that with metadata probes alone and
@@ -263,25 +264,45 @@ async fn mismatch_blob_gaps(
                 || records
                     .first()
                     .is_some_and(|(key, _)| key.as_str() != stripped));
+        // The copies the apply loop gates per copy: gem and PyPI patch
+        // every copy, each against its own representative check; Maven
+        // gates (and patches) only the first.
+        let gate_copies: &[PathBuf] = if matches!(
+            Ecosystem::from_purl(purl),
+            Some(Ecosystem::Gem | Ecosystem::Pypi)
+        ) {
+            pkg_paths.as_slice()
+        } else {
+            std::slice::from_ref(first_path)
+        };
         for (_, record) in records {
             if !can_queue(record) {
                 continue;
             }
-            if gated {
-                if let Some((file_name, file_info)) = representative_file(&record.files) {
-                    let status = verify_file_patch(first_path, file_name, file_info)
-                        .await
-                        .status;
-                    if !variant_matches_installed(Some(&status)) {
-                        continue;
+            // Copies this variant is attempted on: a copy whose installed
+            // distribution is another variant (two envs can hold different
+            // wheels of one release) is skipped there by the apply loop.
+            let probe_copies: Vec<&PathBuf> = match representative_file(&record.files) {
+                Some((file_name, file_info)) if gated => {
+                    let mut matched = Vec::new();
+                    for copy in gate_copies {
+                        let status = verify_file_patch(copy, file_name, file_info).await.status;
+                        if variant_matches_installed(Some(&status)) {
+                            matched.push(copy);
+                        }
                     }
+                    matched
                 }
+                _ => pkg_paths.iter().collect(),
+            };
+            if probe_copies.is_empty() {
+                continue;
             }
             for (file_name, info) in &record.files {
                 if info.before_hash.is_empty() || !missing.contains(&info.after_hash) {
                     continue;
                 }
-                for pkg_path in pkg_paths {
+                for pkg_path in &probe_copies {
                     let verify = verify_file_patch(pkg_path, file_name, info).await;
                     if verify.status == VerifyStatus::HashMismatch {
                         needed.insert(info.after_hash.clone());
@@ -627,6 +648,25 @@ pub(crate) fn variant_matches_installed(first_file_status: Option<&VerifyStatus>
     }
 }
 
+/// `paths` in order with every path that resolves to an already-listed
+/// directory dropped: two discovered site-packages paths can name ONE
+/// directory (a `lib64 -> lib` symlink, a symlinked venv), and patching it
+/// twice would report the second pass `already_patched`. A path that can't
+/// be canonicalized is kept as-is.
+async fn distinct_install_dirs(paths: &[PathBuf]) -> Vec<PathBuf> {
+    let mut seen: HashSet<PathBuf> = HashSet::new();
+    let mut out = Vec::with_capacity(paths.len());
+    for path in paths {
+        let key = tokio::fs::canonicalize(path)
+            .await
+            .unwrap_or_else(|_| path.clone());
+        if seen.insert(key) {
+            out.push(path.clone());
+        }
+    }
+    out
+}
+
 /// The file whose verify status decides whether a release variant
 /// describes the installed distribution (fed to
 /// [`variant_matches_installed`]).
@@ -1881,22 +1921,31 @@ async fn apply_patches_inner(
                 continue;
             }
 
-            // Patch EVERY coexisting gem store copy (the npm multi-copy
-            // precedent): leaving the other store pristine is a silent
-            // false "applied" for whichever bundler loads it, and the
-            // per-copy results below make the JSON summary count each
-            // patched copy — the signal a second copy exists. PyPI/Maven
-            // keep the one-representative contract: their crawlers resolve
-            // one install dir per version, and a second path can only
-            // alias the same logical install (re-patching it would produce
-            // the spurious `already_patched` double-patch the nuget
-            // first-wins restoration fixed).
-            let copy_paths: &[PathBuf] =
-                if matches!(Ecosystem::from_purl(purl), Some(Ecosystem::Gem)) {
-                    pkg_paths.as_slice()
-                } else {
-                    std::slice::from_ref(pkg_path)
-                };
+            // Patch EVERY coexisting gem store copy and every PyPI
+            // site-packages copy (the npm multi-copy precedent): leaving
+            // the other copy pristine is a silent false "applied" for
+            // whichever bundler / interpreter loads it, and the per-copy
+            // results below make the JSON summary count each patched copy
+            // — the signal a second copy exists. The Python crawler
+            // resolves one release in several candidate envs when it
+            // can't tell which one the project's tool runs (a Pipenv
+            // WORKON_HOME venv beside `./.venv`, #529) or which one
+            // `sys.path` shadows (the user site beside a system dir in
+            // global scope, #501), and rollback already restores every
+            // copy. A PyPI path that only ALIASES another (a symlinked
+            // site-packages) is collapsed by canonical path, so one
+            // install is never patched twice. Maven keeps the
+            // one-representative contract: its crawler resolves one
+            // install dir per version.
+            let pypi_copies: Vec<PathBuf>;
+            let copy_paths: &[PathBuf] = match Ecosystem::from_purl(purl) {
+                Some(Ecosystem::Gem) => pkg_paths.as_slice(),
+                Some(Ecosystem::Pypi) => {
+                    pypi_copies = distinct_install_dirs(pkg_paths).await;
+                    pypi_copies.as_slice()
+                }
+                _ => std::slice::from_ref(pkg_path),
+            };
 
             // Copy CLASS decides FAILURE semantics (never write scope —
             // patching a shared home's vulnerable copy is fine when it
@@ -2880,6 +2929,90 @@ mod tests {
         );
     }
 
+    /// Regression (#538 review): two PyPI copies of one release holding
+    /// DIFFERENT wheels. The apply loop gates each variant per copy, so the
+    /// variant installed only in the SECOND copy is attempted there; its
+    /// locally-modified non-representative file needs the full afterHash
+    /// blob. Gating against the first copy alone skipped that variant and
+    /// left the blob unfetched, so warn-and-apply failed on the second copy.
+    #[tokio::test]
+    async fn mismatch_blob_gaps_gates_pypi_variants_per_copy() {
+        use socket_patch_core::hash::git_sha256::compute_git_sha256_from_bytes;
+
+        let dir = tempfile::tempdir().unwrap();
+        // Copy A: the wheel's distribution, pristine.
+        let copy_a = dir.path().join("a");
+        tokio::fs::create_dir_all(&copy_a).await.unwrap();
+        tokio::fs::write(copy_a.join("aaa.py"), b"wheel\n")
+            .await
+            .unwrap();
+        // Copy B: the sdist's distribution, with a locally modified
+        // non-representative file.
+        let copy_b = dir.path().join("b");
+        tokio::fs::create_dir_all(&copy_b).await.unwrap();
+        tokio::fs::write(copy_b.join("aaa.py"), b"sdist\n")
+            .await
+            .unwrap();
+        tokio::fs::write(copy_b.join("zzz.py"), b"locally modified\n")
+            .await
+            .unwrap();
+        let blobs = dir.path().join("blobs");
+        tokio::fs::create_dir_all(&blobs).await.unwrap();
+
+        let mut wheel_files = HashMap::new();
+        wheel_files.insert(
+            "aaa.py".to_string(),
+            PatchFileInfo {
+                before_hash: compute_git_sha256_from_bytes(b"wheel\n"),
+                after_hash: "1".repeat(64),
+            },
+        );
+        let mut manifest = manifest_with_record(
+            "pkg:pypi/foo@1.0.0?artifact_id=foo-1.0.0-py3-none-any.whl",
+            wheel_files,
+        );
+        let mut sdist_files = HashMap::new();
+        sdist_files.insert(
+            "aaa.py".to_string(),
+            PatchFileInfo {
+                before_hash: compute_git_sha256_from_bytes(b"sdist\n"),
+                after_hash: "2".repeat(64),
+            },
+        );
+        sdist_files.insert(
+            "zzz.py".to_string(),
+            PatchFileInfo {
+                before_hash: "3".repeat(64),
+                after_hash: "4".repeat(64),
+            },
+        );
+        manifest.patches.insert(
+            "pkg:pypi/foo@1.0.0?artifact_id=foo-1.0.0.tar.gz".to_string(),
+            PatchRecord {
+                uuid: "22222222-2222-4222-8222-222222222222".to_string(),
+                exported_at: "2024-01-01T00:00:00Z".to_string(),
+                files: sdist_files,
+                vulnerabilities: HashMap::new(),
+                description: "fixture".to_string(),
+                license: "MIT".to_string(),
+                tier: "free".to_string(),
+            },
+        );
+        let mut all_packages = HashMap::new();
+        all_packages.insert(
+            "pkg:pypi/foo@1.0.0".to_string(),
+            vec![copy_a.clone(), copy_b.clone()],
+        );
+
+        let needed =
+            mismatch_blob_gaps(&manifest, &all_packages, &HashSet::new(), &blobs, false).await;
+        assert_eq!(
+            needed,
+            HashSet::from(["4".repeat(64)]),
+            "the second copy's variant must queue its mismatched file's blob"
+        );
+    }
+
     /// A variant with no content-modifying files (only new files) has
     /// nothing to disqualify it: no representative, treated as a match —
     /// the same no-files contract as core's `select_installed_variants`.