Scope Python RPC marketplace responses to the requested distribution (#7512)

* Scope Python RPC marketplace responses to the requested distribution

`discover_recipes()` eagerly walked every `openrewrite.recipes` entry
point on `_get_marketplace()` init and `handle_get_marketplace()`
returned the full singleton, so installing any pip package had its
`RecipeBundle` tagged with every recipe in the process - including the
eight built-in `org.openrewrite.python.*` recipes activated by the
`openrewrite` distribution itself. A visualization-only package like
`moderne-visualizations-misc` ended up "owning" every Python recipe.

Track per-distribution attribution at activation time
(`{distribution_name -> set(recipe_name)}`), have
`handle_install_recipes` return only the rows for the requested
package, and let `handle_get_marketplace` filter by `packageName` when
supplied. `PipRecipeBundleReader` now uses the install response
directly instead of issuing a follow-up `GetMarketplace` call that
would still see the full singleton.

`InstallRecipesResponse` carries the new `recipes` field
backwards-compatibly: older Python servers that don't include it
deserialize to `null`, accessed through `recipesOrEmpty()`.

* Avoid baking the recipe count into prose

The number of built-in `org.openrewrite.python.*` recipes will change
over time. Drop the hard-coded "eight" from the comments and test
docstrings so they don't decay.

* Wrap attribution map in a RecipeAttribution type

`Optional[Dict[str, Set[str]]]` told the reader nothing about what the
keys or values mean, and pushed PEP 503 normalization out to every
call site. Replace it with a small `RecipeAttribution` dataclass that
encapsulates the storage, normalizes distribution names on both
record and lookup, and exposes named `record(...)` / `recipes_for(...)`
methods. Export it from the package since it's now part of the
`discover_recipes` signature.

* Address review feedback on attribution typing + recipesInstalled semantic

- `recipesInstalled` is back to the diff semantic (recipes added by THIS
  call, zero on idempotent reinstalls), matching the field name's promise
  and the prior wire contract. The new `recipes` field carries the
  cumulative per-distribution row list separately. Audit of other
  language servers (Java in-process, C# LanguageServer, JS, Go) confirms
  Python is the sole InstallRecipes producer and no consumer reads
  `recipesInstalled` today, so the diff semantic is the safe default.
- Drop `@dataclass` on `RecipeAttribution`. The auto-generated
  `__init__(_by_package=...)` exposed the internal map through the
  constructor signature and we don't use the auto `__eq__`/`__repr__`.
  Plain class with `__init__(self) -> None` keeps the storage truly
  private.
- Replace bare `Dict[str, Set[str]]` with NewType-typed
  `Dict[NormalizedDistName, Set[RecipeName]]`. `NormalizedDistName` can
  only be built via `_normalize_package_name`; `RecipeName` propagates
  through `recipe_name_set`, the attribution API, and
  `_collect_marketplace_rows`'s filter parameter. Zero runtime cost,
  type-checker discipline at every site.

Author: jkschneider

rewrite-python/rewrite/src/rewrite/__init__.py

@@ -9,7 +9,13 @@
 
 from .category import CategoryDescriptor, LOWEST_PRECEDENCE, DEFAULT_PRECEDENCE, HIGHEST_PRECEDENCE
 from .decorators import categorize, get_recipe_category
-from .discovery import discover_recipes, discover_decorated_recipes_in_module
+from .discovery import (
+    discover_recipes,
+    discover_decorated_recipes_in_module,
+    RecipeAttribution,
+    RecipeName,
+    NormalizedDistName,
+)
 from .execution import ExecutionContext, DelegatingExecutionContext, InMemoryExecutionContext, \
     RecipeRunException, LargeSourceSet, InMemoryLargeSourceSet, Result
 from .marketplace import RecipeMarketplace, Python
@@ -68,6 +74,9 @@
     'get_recipe_category',
     'discover_recipes',
     'discover_decorated_recipes_in_module',
+    'RecipeAttribution',
+    'RecipeName',
+    'NormalizedDistName',
     'activate',
 
     # Markers

rewrite-python/rewrite/src/rewrite/discovery.py

@@ -18,24 +18,81 @@
 
 import inspect
 from importlib.metadata import entry_points
-from typing import List, Tuple, Type
+from typing import Dict, List, NewType, Optional, Set, Tuple, Type
 
 from rewrite.category import CategoryDescriptor
 from rewrite.decorators import get_recipe_category
 from rewrite.marketplace import RecipeMarketplace
 from rewrite.recipe import Recipe
 
 
-def discover_recipes() -> RecipeMarketplace:
+# A recipe's fully qualified name (e.g., ``org.openrewrite.python.RemovePass``).
+# NewType is a static-only distinction — at runtime these are plain ``str`` —
+# but it lets type checkers refuse a bare string where a recipe-name set is
+# expected, which catches the kind of "did I just lookup the wrong key" bug
+# that motivated this module's per-distribution attribution in the first place.
+RecipeName = NewType("RecipeName", str)
+
+# A python distribution name in its PEP 503 normalized form (hyphens,
+# underscores, dots all folded to ``_`` and lowercased). Used as the
+# attribution map key. Construct only via :func:`_normalize_package_name`.
+NormalizedDistName = NewType("NormalizedDistName", str)
+
+
+class RecipeAttribution:
+    """Tracks which distribution's entry point activated which recipes.
+
+    Used to answer "which recipes came from package X?" so callers can scope
+    a marketplace response to a specific distribution instead of returning
+    the whole singleton. Distribution names are PEP 503 normalized (hyphen,
+    underscore, and case folded together) on both write and read, so callers
+    can use any common spelling.
+    """
+
+    def __init__(self) -> None:
+        self._by_package: Dict[NormalizedDistName, Set[RecipeName]] = {}
+
+    def record(self, distribution_name: str, recipe_names: Set[RecipeName]) -> None:
+        """Attribute ``recipe_names`` to ``distribution_name``.
+
+        No-op when ``recipe_names`` is empty; multiple calls for the same
+        distribution accumulate.
+        """
+        if not recipe_names:
+            return
+        key = _normalize_package_name(distribution_name)
+        self._by_package.setdefault(key, set()).update(recipe_names)
+
+    def recipes_for(self, distribution_name: str) -> Set[RecipeName]:
+        """Return the (possibly empty) set of recipe names attributed to a
+        distribution. The returned set is a snapshot — mutating it does not
+        change the attribution.
+        """
+        return set(self._by_package.get(_normalize_package_name(distribution_name), ()))
+
+
+def discover_recipes(
+    marketplace: Optional[RecipeMarketplace] = None,
+    attribution: Optional[RecipeAttribution] = None,
+) -> RecipeMarketplace:
     """
     Discover all recipes from installed packages via entry points.
 
     Looks for packages that declare entry points under
     [project.entry-points."openrewrite.recipes"] and calls their
     activate function to install recipes into the marketplace.
 
+    Args:
+        marketplace: Optional existing marketplace to install into; a new one
+            is created when None.
+        attribution: Optional sink that records which distribution activated
+            which recipes. When supplied, each entry point's contribution is
+            recorded against its distribution name. Recipes already in the
+            marketplace before this call (e.g., deduped by an earlier entry
+            point) are not re-recorded for the current entry point.
+
     Returns:
-        A RecipeMarketplace containing all discovered recipes.
+        The marketplace containing all discovered recipes.
 
     Example pyproject.toml entry point:
         [project.entry-points."openrewrite.recipes"]
@@ -45,23 +102,46 @@ def discover_recipes() -> RecipeMarketplace:
         def activate(marketplace: RecipeMarketplace) -> None:
             marketplace.install(MyRecipe, [Python, Cleanup])
     """
-    marketplace = RecipeMarketplace()
+    if marketplace is None:
+        marketplace = RecipeMarketplace()
 
     # Python 3.10+ uses select parameter via SelectableGroups
     eps = entry_points(group="openrewrite.recipes")
 
     for ep in eps:
         try:
             module = ep.load()
-            if hasattr(module, "activate") and callable(module.activate):
+            if not hasattr(module, "activate") or not callable(module.activate):
+                continue
+            if attribution is None:
                 module.activate(marketplace)
+                continue
+            dist_name = ep.dist.name if ep.dist is not None else None
+            before = recipe_name_set(marketplace)
+            module.activate(marketplace)
+            if dist_name:
+                attribution.record(dist_name, recipe_name_set(marketplace) - before)
         except Exception:
             # Log or handle the error - for now, skip failed activations
             pass
 
     return marketplace
 
 
+def recipe_name_set(marketplace: RecipeMarketplace) -> Set[RecipeName]:
+    """Snapshot the set of recipe names currently in the marketplace."""
+    return {RecipeName(r.name) for r in marketplace.all_recipes()}
+
+
+def _normalize_package_name(name: str) -> NormalizedDistName:
+    """Normalize a python distribution name for attribution lookup.
+
+    PyPI/pip treats hyphens, underscores, and case as equivalent when
+    resolving distribution identity.
+    """
+    return NormalizedDistName(name.replace("-", "_").replace(".", "_").lower())
+
+
 def discover_decorated_recipes_in_module(
     module,
 ) -> List[Tuple[Type[Recipe], List[CategoryDescriptor]]]:

rewrite-python/rewrite/src/rewrite/rpc/server.py

@@ -33,9 +33,11 @@
 import threading
 
 from pathlib import Path
-from typing import Dict, Any, Optional, List, Callable
+from typing import Dict, Any, Optional, List, Callable, Set
 from uuid import uuid4
 
+from rewrite.discovery import RecipeAttribution, RecipeName
+
 # Deeply nested LST nodes (e.g., 256 implicitly concatenated strings) can
 # overflow the default recursion limit (1000) during RPC serialization.
 sys.setrecursionlimit(sys.getrecursionlimit() * 2)
@@ -584,21 +586,43 @@ def handle_reset(params: dict) -> bool:
 # Global marketplace instance (lazily initialized)
 _marketplace = None
 
+# Tracks which distribution's entry point activated which recipes. Used by
+# handle_install_recipes to scope its response to the requested distribution
+# and by handle_get_marketplace to filter the singleton marketplace down to
+# a single package.
+_attribution = RecipeAttribution()
+
 
 def _get_marketplace():
-    """Get or create the global marketplace instance."""
+    """Get or create the global marketplace instance.
+
+    Discovery populates ``_attribution`` so each recipe is attributed to the
+    distribution whose entry point activated it. Without that attribution, a
+    later GetMarketplace or InstallRecipes for package X would incorrectly
+    return every recipe in the singleton, including the built-in
+    ``org.openrewrite.python.*`` recipes activated by the ``openrewrite``
+    distribution itself.
+    """
     global _marketplace
     if _marketplace is None:
-        from rewrite.discovery import discover_recipes
+        from rewrite.discovery import discover_recipes, recipe_name_set
         from rewrite.marketplace import RecipeMarketplace
         from rewrite import activate
 
-        # First try to discover from installed packages
-        _marketplace = discover_recipes()
+        _marketplace = RecipeMarketplace()
 
-        # Also activate local recipes (in case package isn't installed)
-        # This ensures recipes work during development
+        # Discover from installed packages, tracking which distribution
+        # contributed each recipe.
+        discover_recipes(marketplace=_marketplace, attribution=_attribution)
+
+        # Also activate local recipes (in case the openrewrite distribution
+        # isn't pip-installed, e.g., when running from source). When it is
+        # installed, discovery already covered these and install() will dedupe
+        # by name; attribute the source-mode additions to "openrewrite" so
+        # they're returned by GetMarketplace/InstallRecipes for that package.
+        before = recipe_name_set(_marketplace)
         activate(_marketplace)
+        _attribution.record("openrewrite", recipe_name_set(_marketplace) - before)
 
     return _marketplace
 
@@ -616,16 +640,24 @@ def handle_install_recipes(params: dict) -> dict:
             - 'recipes': {'packageName': str, 'version': str|None} - A package spec
 
     Returns:
-        Dict with 'recipesInstalled' count and 'version' (if resolved)
+        Dict with:
+            - 'recipesInstalled': count of recipes added to the marketplace by
+              this call (zero on idempotent reinstalls).
+            - 'version': resolved version (if known).
+            - 'recipes': cumulative list of {descriptor, categoryPaths} rows
+              for recipes attributed to this distribution. Stable across
+              reinstalls; this is what the caller binds to its bundle.
     """
     import importlib
     import importlib.util
+    from rewrite.discovery import recipe_name_set
 
     marketplace = _get_marketplace()
-    before_count = len(list(marketplace.all_recipes()))
 
     recipes = params.get('recipes')
     installed_version = None
+    package_name: Optional[str] = None
+    recipes_added = 0
 
     if isinstance(recipes, str):
         # Local file path - package should already be installed by caller
@@ -636,7 +668,11 @@ def handle_install_recipes(params: dict) -> dict:
         # For local paths, we look for the package name from setup.py/pyproject.toml
         package_name = _find_package_name(local_path)
         if package_name:
+            before = recipe_name_set(marketplace)
             _import_and_activate_package(package_name, marketplace, local_path)
+            added = recipe_name_set(marketplace) - before
+            _attribution.record(package_name, added)
+            recipes_added = len(added)
 
     elif isinstance(recipes, dict):
         # Package spec with name and optional version - package should already be installed
@@ -655,17 +691,31 @@ def handle_install_recipes(params: dict) -> dict:
         except Exception:
             pass
 
+        before = recipe_name_set(marketplace)
         _import_and_activate_package(package_name, marketplace)
+        added = recipe_name_set(marketplace) - before
+        _attribution.record(package_name, added)
+        recipes_added = len(added)
     else:
         raise ValueError(f"Invalid recipes parameter: {recipes}")
 
-    after_count = len(list(marketplace.all_recipes()))
-    recipes_installed = after_count - before_count
-
-    logger.info(f"InstallRecipes: installed {recipes_installed} recipes")
+    package_rows: List[dict] = []
+    if package_name:
+        # recipes_for returns an empty set (not None) when the package activated
+        # nothing, so the filter scopes the result to "this package's recipes"
+        # (zero of them) instead of falling back to "no filter" and returning
+        # everything.
+        package_rows = _collect_marketplace_rows(
+            marketplace, recipe_filter=_attribution.recipes_for(package_name)
+        )
+
+    logger.info(
+        f"InstallRecipes: {recipes_added} new, {len(package_rows)} cumulative for {package_name}"
+    )
     return {
-        'recipesInstalled': recipes_installed,
-        'version': installed_version
+        'recipesInstalled': recipes_added,
+        'version': installed_version,
+        'recipes': package_rows,
     }
 
 
@@ -828,22 +878,48 @@ def normalize(name: str) -> str:
 def handle_get_marketplace(params: dict) -> List[dict]:
     """Handle a GetMarketplace RPC request.
 
-    Returns all recipes organized by category in a format compatible with Java.
+    Returns recipes organized by category in a format compatible with Java.
+
+    Args:
+        params: Optional dict with a 'packageName' key. When supplied, the
+            response is filtered to recipes attributed to that distribution,
+            so callers requesting a specific bundle don't get every recipe
+            in the singleton marketplace.
 
     Returns:
         List of dicts with 'descriptor' and 'categoryPaths' for each recipe.
     """
-    from dataclasses import asdict
-
     marketplace = _get_marketplace()
+
+    recipe_filter: Optional[Set[RecipeName]] = None
+    if isinstance(params, dict):
+        package_name = params.get('packageName')
+        if isinstance(package_name, str) and package_name:
+            recipe_filter = _attribution.recipes_for(package_name)
+
+    rows = _collect_marketplace_rows(marketplace, recipe_filter=recipe_filter)
+    logger.info(f"GetMarketplace: returning {len(rows)} recipes")
+    return rows
+
+
+def _collect_marketplace_rows(
+    marketplace,
+    recipe_filter: Optional[Set[RecipeName]] = None,
+) -> List[dict]:
+    """Walk the marketplace and return recipe rows in GetMarketplaceResponse shape.
+
+    A recipe that appears in multiple categories produces one row whose
+    ``categoryPaths`` lists each path. When ``recipe_filter`` is provided,
+    recipes whose name isn't in the set are skipped.
+    """
     rows: List[dict] = []
 
-    def collect_recipes(category, category_path: List[dict]):
-        """Recursively collect recipes from a category and its subcategories."""
+    def collect(category, category_path: List[dict]) -> None:
         current_path = [*category_path, _category_descriptor_to_dict(category.descriptor)]
 
-        for recipe_name, (recipe_desc, _recipe_class) in category.recipes.items():
-            # Check if we already have this recipe (it can appear in multiple categories)
+        for _recipe_name, (recipe_desc, _recipe_class) in category.recipes.items():
+            if recipe_filter is not None and recipe_desc.name not in recipe_filter:
+                continue
             existing = next((r for r in rows if r['descriptor']['name'] == recipe_desc.name), None)
             if existing:
                 existing['categoryPaths'].append(current_path)
@@ -854,13 +930,11 @@ def collect_recipes(category, category_path: List[dict]):
                 })
 
         for subcategory in category.categories:
-            collect_recipes(subcategory, current_path)
+            collect(subcategory, current_path)
 
-    # Start from the root's children (skip the root itself)
     for category in marketplace.categories():
-        collect_recipes(category, [])
+        collect(category, [])
 
-    logger.info(f"GetMarketplace: returning {len(rows)} recipes")
     return rows