RPC: Fix getObject() desync bugs (#6835)

* RPC: Fix getObject() using wrong baseline for receiving diffs

getObject() used localObjects as the diff baseline when receiving from
the remote peer. But the remote computes diffs against the last synced
state (remoteObjects), not the local state. When the local side modifies
a tree (e.g., via a local recipe) before calling getObject(), the two
baselines diverge, producing a hybrid tree that corrupts remoteObjects
and causes subsequent transfers to fail with IndexError/desync.

Fix: use remoteObjects (the last synced state) as the baseline, matching
what the remote peer uses. Python's get_object_from_java() already did
this correctly; Java and TypeScript had the same bug.

* RPC: Reset remote object tracking on getObject() failure to prevent state desync

When getObject() fails mid-deserialization (e.g., ClassCastException), the
sender has already updated its tracking of what the receiver has, but the
receiver never applied the change. This causes all subsequent RPC interactions
for that object to compute diffs against the wrong baseline, leading to
"Expected positions array" and similar desync errors.

Fix by removing remoteObjects[id] when receive() throws, so the next
interaction sends a full ADD (no delta), re-synchronizing both sides.

Also fix Python's handle_visit() to always fetch the tree from Java via
get_object_from_java(), matching the JS implementation. The previous code
used a stale local_objects cache, which meant Python could operate on an
outdated tree version after Java-side modifications or error recovery.

* 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-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java

@@ -461,16 +461,27 @@ List<String> getCursorIds(@Nullable Cursor cursor) {
 
     @VisibleForTesting
     public <T> T getObject(String id, @Nullable String sourceFileType) {
-        // Check if we have a cached version of this object
-        Object localObject = localObjects.get(id);
+        // Use the last synced state as the baseline for receiving diffs.
+        // This must match what the remote used as its baseline when computing the diff.
+        // Using localObjects here would be wrong if Java modified the tree locally
+        // (e.g., via a Java-side recipe) since the remote doesn't know about those changes.
+        Object before = remoteObjects.get(id);
 
         RpcReceiveQueue q = new RpcReceiveQueue(
                 remoteRefs,
                 () -> send("GetObject", new GetObject(id, sourceFileType), GetObjectResponse.class),
                 sourceFileType,
                 log.get()
         );
-        Object remoteObject = q.receive(localObject, null);
+        Object remoteObject;
+        try {
+            remoteObject = q.receive(before, null);
+        } catch (Exception e) {
+            // Reset our tracking of the remote state so the next interaction
+            // forces a full object sync (ADD) instead of a delta (CHANGE).
+            remoteObjects.remove(id);
+            throw e;
+        }
         if (q.take().getState() != END_OF_OBJECT) {
             throw new IllegalStateException("Expected END_OF_OBJECT");
         }

rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java

@@ -25,6 +25,7 @@
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.openrewrite.*;
+import org.openrewrite.marker.Markers;
 import org.openrewrite.config.Environment;
 import org.openrewrite.config.OptionDescriptor;
 import org.openrewrite.config.RecipeDescriptor;
@@ -38,8 +39,10 @@
 import java.io.IOException;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
+import java.nio.file.Path;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 
 import static java.util.Objects.requireNonNull;
@@ -82,6 +85,57 @@ void after() {
         server.shutdown();
     }
 
+    /**
+     * Verifies that getObject() uses remoteObjects (last synced state) as the
+     * diff baseline, not localObjects. When the local side modifies a tree
+     * (e.g., via a local recipe) before calling getObject(), the localObjects
+     * entry diverges from what the remote used as its baseline. Using
+     * localObjects would cause NO_CHANGE fields to retain the local
+     * modification rather than the synced value.
+     * <p>
+     * The bug only manifests for object fields (identity-compared via ==)
+     * where the sender emits NO_CHANGE. Here, the server changes only text
+     * (a value field), so Markers — an object field preserved by reference
+     * in withText() — is sent as NO_CHANGE. If the client locally created
+     * a different Markers object, the wrong baseline would retain it.
+     */
+    @Test
+    void getObjectUsesRemoteObjectsAsBaseline() {
+        PlainText original = PlainText.builder()
+          .sourcePath(Path.of("test.txt"))
+          .text("Hello")
+          .build();
+
+        String id = original.getId().toString();
+        String sourceFileType = PlainText.class.getName();
+
+        // Server has the tree; client fetches it → both synced
+        server.localObjects.put(id, original);
+        PlainText synced = client.getObject(id, sourceFileType);
+        UUID syncedMarkersId = synced.getMarkers().getId();
+
+        // Server modifies text only — Markers reference stays the same
+        // (original.withText() preserves the Markers by reference via @With),
+        // so the handler sends NO_CHANGE for the Markers field.
+        // Using original (not synced) ensures server's before == original.getMarkers()
+        // and after == original.withText(...).getMarkers() are the same reference.
+        server.localObjects.put(id, original.withText("Hello World"));
+
+        // Client locally modifies Markers (simulating a local recipe step),
+        // creating a new Markers object with a different ID
+        Markers localMarkers = synced.getMarkers().withId(Tree.randomId());
+        client.localObjects.put(id, synced.withMarkers(localMarkers));
+
+        // Fetch: server sends Markers=NO_CHANGE, text=CHANGE("Hello World").
+        // The receiver should apply NO_CHANGE against remoteObjects (synced
+        // markers), not localObjects (local markers).
+        PlainText result = client.getObject(id, sourceFileType);
+        assertThat(result.getText()).isEqualTo("Hello World");
+        assertThat(result.getMarkers().getId())
+          .describedAs("Markers should come from remoteObjects baseline, not localObjects")
+          .isEqualTo(syncedMarkersId);
+    }
+
     @DocumentExample
     @Test
     void sendReceiveIdempotence() {

rewrite-javascript/rewrite/src/rpc/rewrite-rpc.ts

@@ -124,7 +124,11 @@ export class RewriteRpc {
     }
 
     async getObject<P>(id: string, sourceFileType?: string): Promise<P> {
-        const localObject = this.localObjects.get(id);
+        // Use the last synced state as the baseline for receiving diffs.
+        // This must match what the remote used as its baseline when computing the diff.
+        // Using localObjects here would be wrong if the local side modified the tree
+        // (e.g., via a local recipe) since the remote doesn't know about those changes.
+        const before = this.remoteObjects.get(id);
 
         const q = new RpcReceiveQueue(this.remoteRefs, sourceFileType, () => {
             return this.connection.sendRequest(
@@ -133,7 +137,15 @@ export class RewriteRpc {
             );
         }, this.logger, this.traceGetObject.receive);
 
-        const remoteObject = await q.receive<P>(localObject);
+        let remoteObject: P;
+        try {
+            remoteObject = await q.receive<P>(before as P);
+        } catch (e) {
+            // Reset our tracking of the remote state so the next interaction
+            // forces a full object sync (ADD) instead of a delta (CHANGE).
+            this.remoteObjects.delete(id);
+            throw e;
+        }
 
         const eof = (await q.take());
         if (eof.state !== RpcObjectState.END_OF_OBJECT) {

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

@@ -194,12 +194,22 @@ def pull_batch() -> List[Dict[str, Any]]:
     before = remote_objects.get(obj_id)
 
     # Receive and deserialize the object (applies diffs to before state)
-    obj = receiver.receive(before, q)
-
-    # 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}")
+    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
 
     if obj is not None:
         # Update our understanding of what Java has
@@ -1036,10 +1046,9 @@ def handle_visit(params: dict) -> dict:
         if p_id:
             _execution_contexts[p_id] = ctx
 
-    # Get the tree - fetch from Java if we don't have it locally
-    tree = local_objects.get(tree_id)
-    if tree is None:
-        tree = get_object_from_java(tree_id, source_file_type)
+    # Always fetch the tree from Java to ensure we have the latest version.
+    # Java may have modified the tree (e.g., via a Java-side recipe) since our last sync.
+    tree = get_object_from_java(tree_id, source_file_type)
 
     if tree is None:
         raise ValueError(f"Tree not found: {tree_id}")