RPC: Fix Python get_object_from_java missing END_OF_OBJECT on batch boundary

When serialized data items are an exact multiple of the handler's
batchSize, END_OF_OBJECT ends up alone in a separate batch that
receiver.receive() never pulls. Drain the pending batch after
receive() completes — analogous to Java/JS's explicit q.take().

Add integration test with small batchSize to exercise the fix.

Author: knutwannheden

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

@@ -196,17 +196,21 @@ def pull_batch() -> List[Dict[str, Any]]:
     # Receive and deserialize the object (applies diffs to before state)
     try:
         obj = receiver.receive(before, q)
+
+        # After receive() completes, END_OF_OBJECT may still be pending in a
+        # separate batch (happens when data items are an exact multiple of the
+        # handler's batchSize). Drain it — analogous to Java's explicit
+        # q.take() after receive().
+        if not received_end:
+            pull_batch()
+        if not received_end:
+            raise RuntimeError(f"Did not receive END_OF_OBJECT marker for object {obj_id}")
     except Exception:
         # Reset our tracking of the remote state so the next interaction
         # forces a full object sync (ADD) instead of a delta (CHANGE).
         remote_objects.pop(obj_id, None)
         raise
 
-    # Verify we received the complete object (END_OF_OBJECT was in the final batch)
-    # This matches Java's RewriteRpc.java line 474-475 which explicitly checks for END_OF_OBJECT
-    if not received_end:
-        raise RuntimeError(f"Did not receive END_OF_OBJECT marker for object {obj_id}")
-
     if obj is not None:
         # Update our understanding of what Java has
         remote_objects[obj_id] = obj

rewrite-python/src/integTest/java/org/openrewrite/python/PythonBestPracticesIntegTest.java

@@ -0,0 +1,144 @@
+/*
+ * Copyright 2026 the original author or authors.
+ * <p>
+ * Licensed under the Moderne Source Available License (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://docs.moderne.io/licensing/moderne-source-available-license
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.python;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.api.io.TempDir;
+import org.openrewrite.ExecutionContext;
+import org.openrewrite.InMemoryExecutionContext;
+import org.openrewrite.Recipe;
+import org.openrewrite.RecipeRun;
+import org.openrewrite.Result;
+import org.openrewrite.SourceFile;
+import org.openrewrite.internal.InMemoryLargeSourceSet;
+import org.openrewrite.python.rpc.PythonRewriteRpc;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Integration test that installs recipes from rewrite-migrate-python and
+ * runs PythonBestPractices against real projects.
+ * <p>
+ * This exercises the full RPC lifecycle including recipe installation,
+ * project parsing, recipe execution (Visit + GetObject + Print), and
+ * verifies that the diff-based RPC protocol stays in sync across
+ * multiple recipe steps.
+ */
+class PythonBestPracticesIntegTest {
+
+    static final Path RECIPES_PATH = Path.of("/Users/knut/git/moderneinc/rewrite-migrate-python/.worktrees/nimble-hawk/recipes-python");
+
+    @TempDir
+    Path tempDir;
+
+    @BeforeEach
+    void before() {
+        PythonRewriteRpc.setFactory(PythonRewriteRpc.builder()
+                .log(tempDir.resolve("python-rpc.log"))
+                .traceRpcMessages()
+        );
+    }
+
+    @AfterEach
+    void after() throws IOException {
+        PythonRewriteRpc.shutdownCurrent();
+        PythonRewriteRpc.setFactory(PythonRewriteRpc.builder());
+        Path logFile = tempDir.resolve("python-rpc.log");
+        if (Files.exists(logFile)) {
+            String log = Files.readString(logFile);
+            if (log.contains("ERROR")) {
+                System.out.println("=== Python RPC Log (errors found) ===");
+                System.out.println(log);
+            }
+        }
+    }
+
+    @Test
+    @Timeout(value = 300, unit = TimeUnit.SECONDS)
+    void moreItertools() {
+        runPythonBestPractices(Path.of("/Users/knut/moderne/working-set-python-static-analysis-2/more-itertools/more-itertools"));
+    }
+
+    @Test
+    @Timeout(value = 300, unit = TimeUnit.SECONDS)
+    void marshmallow() {
+        runPythonBestPractices(Path.of("/Users/knut/moderne/working-set-python-static-analysis-2/marshmallow-code/marshmallow"));
+    }
+
+    /**
+     * Runs with a small batch size to force END_OF_OBJECT into a separate batch.
+     * This exercises the fix in Python's get_object_from_java that drains
+     * the END_OF_OBJECT batch after receive() completes.
+     */
+    @Test
+    @Timeout(value = 600, unit = TimeUnit.SECONDS)
+    void marshmallowSmallBatch() {
+        PythonRewriteRpc rpc = PythonRewriteRpc.getOrStart();
+        rpc.batchSize(7);
+        runPythonBestPractices(rpc, Path.of("/Users/knut/moderne/working-set-python-static-analysis-2/marshmallow-code/marshmallow"));
+    }
+
+    private void runPythonBestPractices(Path projectPath) {
+        runPythonBestPractices(PythonRewriteRpc.getOrStart(), projectPath);
+    }
+
+    private void runPythonBestPractices(PythonRewriteRpc rpc, Path projectPath) {
+
+        // Install recipes from local path
+        rpc.installRecipes(RECIPES_PATH.toFile());
+
+        // Prepare the composite recipe
+        Recipe recipe = rpc.prepareRecipe("org.openrewrite.python.cleanup.PythonBestPractices");
+
+        // Parse the project
+        ExecutionContext ctx = new InMemoryExecutionContext(Throwable::printStackTrace);
+        List<SourceFile> sources = rpc.parseProject(projectPath, ctx)
+                .collect(Collectors.toList());
+        assertThat(sources).isNotEmpty();
+        System.out.println("Parsed " + sources.size() + " source files");
+
+        // Run the recipe against each source file individually
+        int changed = 0;
+        for (SourceFile source : sources) {
+            System.out.println("Processing: " + source.getSourcePath());
+            System.out.flush();
+
+            RecipeRun run = recipe.run(
+                    new InMemoryLargeSourceSet(List.of(source)), ctx);
+
+            for (Result result : run.getChangeset().getAllResults()) {
+                String diff = result.diff(null);
+                if (!diff.isEmpty()) {
+                    changed++;
+                    System.out.println("  Changed: " + result.getAfter().getSourcePath());
+                }
+            }
+        }
+
+        System.out.println("Total sources: " + sources.size());
+        System.out.println("Total changed: " + changed);
+    }
+}