RPC: gate BatchVisit by remote getLanguages (#7553)

jkschneider · web-flow · commit becae3be3a7d · 2026-05-03T15:20:08.000-04:00
Without this, every SourceFile in a repository was being shipped to the
remote during a batched RPC recipe run, even types the remote has no
codec for. The remote then crashed deserializing them, and the resulting
exception was attached to the source as a `Markup.Error` marker — falsely
reporting the file as `MODIFY` and causing the Python migration recipes
to "touch" every yaml/Dockerfile/license/README in the workspace.

Three fixes that ship together:

* `RecipeRunCycle.editSource` now consults `currentRpc.getLanguages()`
  before adding a recipe to the batch, mirroring the `isAcceptable` gate
  the non-batched path already had.

* `rewrite/rpc/server.py` validates with `_require_tree(...)` after every
  `get_object_from_java(...)` in `handle_visit` / `handle_batch_visit`.
  When the receiver has no codec for a top-level type it falls back to
  a `{"kind": value_type, ...}` dict; that dict was leaking into the
  visitor framework and crashing with `AttributeError: 'dict' object
  has no attribute 'is_acceptable'`. The guard now raises the same
  canonical "No RPC codec registered" error the ADD path raises, so the
  failure mode is consistent regardless of which RPC message shape the
  remote used.

* `RewriteRpc`'s `GetLanguages` handler used to advertise only
  `PlainText`, `Json$Document`, `J$CompilationUnit`, and
  `JS$CompilationUnit`. Java actually has receivers for every standard
  format (Yaml, Xml, Hcl, Properties, Toml, Protobuf, Docker) and every
  language compilation unit (J, JS, K, G, Cs, Py, Go, S). The handler
  now lists all of them, each gated by `ifOnClasspath` so the advertised
  set still narrows to whichever rewrite-* modules are actually loaded.
diff --git a/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java b/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java
@@ -176,11 +176,29 @@ protected Boolean handle(TraceGetObject request) {
         jsonRpc.rpc("GetLanguages", new JsonRpcMethod<Void>() {
             @Override
             protected Object handle(Void noParams) {
+                // Advertise every SourceFile type Java has a receiver for so that
+                // remote peers know they can hand any of these off to Java for
+                // visiting. Each entry is gated by ifOnClasspath so the actual
+                // set narrows to whichever rewrite-* modules are loaded in this
+                // process.
                 return Stream.of(
                         ifOnClasspath("org.openrewrite.text.PlainText"),
                         ifOnClasspath("org.openrewrite.json.tree.Json$Document"),
+                        ifOnClasspath("org.openrewrite.yaml.tree.Yaml$Documents"),
+                        ifOnClasspath("org.openrewrite.xml.tree.Xml$Document"),
+                        ifOnClasspath("org.openrewrite.hcl.tree.Hcl$ConfigFile"),
+                        ifOnClasspath("org.openrewrite.properties.tree.Properties$File"),
+                        ifOnClasspath("org.openrewrite.toml.tree.Toml$Document"),
+                        ifOnClasspath("org.openrewrite.protobuf.tree.Proto$Document"),
+                        ifOnClasspath("org.openrewrite.docker.tree.Docker$File"),
                         ifOnClasspath("org.openrewrite.java.tree.J$CompilationUnit"),
-                        ifOnClasspath("org.openrewrite.javascript.tree.JS$CompilationUnit")
+                        ifOnClasspath("org.openrewrite.javascript.tree.JS$CompilationUnit"),
+                        ifOnClasspath("org.openrewrite.kotlin.tree.K$CompilationUnit"),
+                        ifOnClasspath("org.openrewrite.groovy.tree.G$CompilationUnit"),
+                        ifOnClasspath("org.openrewrite.csharp.tree.Cs$CompilationUnit"),
+                        ifOnClasspath("org.openrewrite.python.tree.Py$CompilationUnit"),
+                        ifOnClasspath("org.openrewrite.golang.tree.Go$CompilationUnit"),
+                        ifOnClasspath("org.openrewrite.scala.tree.S$CompilationUnit")
                 ).filter(Objects::nonNull).toArray(String[]::new);
             }
 
diff --git a/rewrite-core/src/main/java/org/openrewrite/scheduling/RecipeRunCycle.java b/rewrite-core/src/main/java/org/openrewrite/scheduling/RecipeRunCycle.java
@@ -327,14 +327,20 @@ void clear() {
                 }
 
                 if (isInBatch) {
-                    // Batch path: accumulate visitor names instead of executing
-                    RpcRecipe rpcRecipe = (RpcRecipe) recipe;
-                    batch.items.add(new BatchVisit.BatchVisitItem(rpcRecipe.getEditVisitor(), null));
-                    batch.recipeStacks.add(recipeStack);
-                    if (batch.originalBeforeBatch == null) {
-                        batch.originalBeforeBatch = src;
+                    // Batch path: accumulate visitor names instead of executing.
+                    // Only batch if the remote actually has a codec for this source file type;
+                    // otherwise the BatchVisit RPC fails on the remote side and the resulting
+                    // exception gets attached to the source as a Markup$Error, falsely marking
+                    // the file as modified.
+                    if (currentRpc.getLanguages().contains(src.getClass().getName())) {
+                        RpcRecipe rpcRecipe = (RpcRecipe) recipe;
+                        batch.items.add(new BatchVisit.BatchVisitItem(rpcRecipe.getEditVisitor(), null));
+                        batch.recipeStacks.add(recipeStack);
+                        if (batch.originalBeforeBatch == null) {
+                            batch.originalBeforeBatch = src;
+                        }
+                        batch.rpc = currentRpc;
                     }
-                    batch.rpc = currentRpc;
 
                     // If this is the last recipe in the batch, flush now
                     if (nextRpc != currentRpc) {
diff --git a/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java
@@ -451,6 +451,35 @@ void batchVisitExceptionProducesErrorMarker() {
           .isPresent();
     }
 
+    /**
+     * Source files whose type the remote does not advertise via GetLanguages must
+     * not be added to a BatchVisit. Without this gate the remote's receiver crashes
+     * on the unknown type and the resulting exception is attached to the file as a
+     * Markup.Error marker, falsely reporting the file as modified.
+     */
+    @Test
+    void batchSkipsSourceFilesNotInRemoteLanguages() {
+        // given
+        // Two consecutive same-RPC recipes are required for the batch path to kick in
+        Recipe r1 = client.prepareRecipe("org.openrewrite.text.ChangeText", Map.of("toText", "step1"));
+        Recipe r2 = client.prepareRecipe("org.openrewrite.text.ChangeText", Map.of("toText", "step2"));
+
+        // Quark is not in the hardcoded GetLanguages list, so the batch path must skip it.
+        var quark = new org.openrewrite.quark.Quark(UUID.randomUUID(), Path.of("unknown.bin"), Markers.EMPTY, null, null);
+
+        var ctx = new InMemoryExecutionContext();
+        RecipeRun run = new RecipeScheduler().scheduleRun(
+          new CompositeRecipe(List.of(r1, r2)),
+          new InMemoryLargeSourceSet(List.of(quark)), ctx, 1, 1);
+
+        // then
+        assertThat(run.getChangeset().getAllResults())
+          .describedAs("Quark of an unsupported language type should pass through untouched")
+          .allSatisfy(result -> assertThat(result.getAfter().getMarkers().findFirst(Markup.Error.class))
+            .describedAs("No Markup.Error should be attached for an unsupported source type")
+            .isEmpty());
+    }
+
     @Test
     void getCursor() {
         var parent = new Cursor(null, Cursor.ROOT_VALUE);
diff --git a/rewrite-python/rewrite/src/rewrite/rpc/server.py b/rewrite-python/rewrite/src/rewrite/rpc/server.py
@@ -137,6 +137,29 @@ def send_request(method: str, params: dict, timeout_seconds: float = 30.0) -> An
     return response.get('result')
 
 
+def _require_tree(tree: Any, source_file_type: Optional[str]) -> Any:
+    """Validate that ``tree`` is a real Tree, not a generic dict fallback.
+
+    When the receiver encounters a ``value_type`` it has no codec for, it
+    falls back to returning a ``{'kind': value_type, ...}`` dict (see
+    ``RpcReceiveQueue._new_obj`` / ``_do_change``). That fallback is
+    appropriate for nested fragments the visitor framework never inspects
+    directly, but a top-level SourceFile *must* be a Tree — otherwise the
+    visitor crashes with a confusing ``AttributeError: 'dict' object has
+    no attribute 'is_acceptable'``. Raise the same "No RPC codec" error
+    that the ADD path raises so the failure mode is consistent regardless
+    of which RPC message shape Java used.
+    """
+    from rewrite import Tree
+    if isinstance(tree, Tree):
+        return tree
+    raise RuntimeError(
+        f"No RPC codec registered on the Python side for '{source_file_type}'. "
+        "The remote side has a codec and sent property messages that will not be consumed, "
+        "causing RPC queue desynchronization."
+    )
+
+
 def get_object_from_java(obj_id: str, source_file_type: Optional[str] = None) -> Any:
     """Fetch an object from Java and deserialize it using PythonRpcReceiver.
 
@@ -1199,6 +1222,8 @@ def handle_visit(params: dict) -> dict:
     if tree is None:
         raise ValueError(f"Tree not found: {tree_id}")
 
+    tree = _require_tree(tree, source_file_type)
+
     # Instantiate the visitor
     visitor = _instantiate_visitor(visitor_name, ctx)
 
@@ -1272,6 +1297,8 @@ def handle_batch_visit(params: dict) -> dict:
     if tree is None:
         raise ValueError(f"Tree not found: {tree_id}")
 
+    tree = _require_tree(tree, source_file_type)
+
     from rewrite.visitor import Cursor
     from rewrite.markers import SearchResult
     cursor = Cursor(None, Cursor.ROOT_VALUE)
diff --git a/rewrite-python/rewrite/tests/rpc/test_server.py b/rewrite-python/rewrite/tests/rpc/test_server.py
@@ -1,3 +1,24 @@
+def test_require_tree_rejects_dict_fallback():
+    """A top-level SourceFile must be a Tree. The receiver returns a generic
+    ``{'kind': value_type, ...}`` dict for value_types it has no codec for —
+    that fallback is fine for nested fragments but lets unknown source files
+    leak into the visitor framework, where they crash with
+    ``AttributeError: 'dict' object has no attribute 'is_acceptable'``.
+    Verify the guard converts that into the same "No RPC codec" error the
+    ADD path raises so the failure mode is consistent."""
+    import pytest
+    from rewrite.rpc.server import _require_tree
+
+    with pytest.raises(RuntimeError, match="No RPC codec registered"):
+        _require_tree(
+            {"kind": "org.openrewrite.docker.tree.Docker$Document"},
+            "org.openrewrite.docker.tree.Docker$Document",
+        )
+
+    with pytest.raises(RuntimeError, match="No RPC codec registered"):
+        _require_tree(None, "org.openrewrite.text.PlainText")
+
+
 def test_handle_parse_preserves_empty_text(tmp_path, monkeypatch):
     import rewrite.rpc.server as server
 
