From b00159a641fb4887a7e8d8e1f61d26410fd061e5 Mon Sep 17 00:00:00 2001
From: Dominic Letz <dominic@diode.io>
Date: Fri, 4 Sep 2026 17:43:26 +0200
Subject: [PATCH] Fix stale NodeProxy reconnect state

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 lib/remote_chain/node_proxy.ex        | 15 +++++++-
 test/remote_chain/node_proxy_test.exs | 53 +++++++++++++++++++++++----
 2 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/lib/remote_chain/node_proxy.ex b/lib/remote_chain/node_proxy.ex
index 7d3c0ef..3fd8f3c 100644
--- a/lib/remote_chain/node_proxy.ex
+++ b/lib/remote_chain/node_proxy.ex
@@ -431,7 +431,8 @@ defmodule RemoteChain.NodeProxy do
            connections: connections,
            requests: requests,
            fallback: fallback,
-           fallback_url: fallback_url
+           fallback_url: fallback_url,
+           lastblocks: lastblocks
          },
          down_pid
        ) do
@@ -446,7 +447,18 @@ defmodule RemoteChain.NodeProxy do
       end)
       |> Map.new()
 
+    removed_urls =
+      for {url, pid} <- connections, pid == down_pid, do: url
+
+    removed_urls =
+      if fallback == down_pid and not is_nil(fallback_url) do
+        [fallback_url | removed_urls]
+      else
+        removed_urls
+      end
+
     new_connections = Enum.filter(connections, fn {_, pid} -> pid != down_pid end) |> Map.new()
+    lastblocks = Map.drop(lastblocks, removed_urls)
 
     {new_fallback, new_fallback_url} =
       if fallback == down_pid, do: {nil, nil}, else: {fallback, fallback_url}
@@ -455,6 +467,7 @@ defmodule RemoteChain.NodeProxy do
       state
       | connections: new_connections,
         requests: requests,
+        lastblocks: lastblocks,
         fallback: new_fallback,
         fallback_url: new_fallback_url
     }
diff --git a/test/remote_chain/node_proxy_test.exs b/test/remote_chain/node_proxy_test.exs
index 0cb41cb..dc8bbbb 100644
--- a/test/remote_chain/node_proxy_test.exs
+++ b/test/remote_chain/node_proxy_test.exs
@@ -591,6 +591,7 @@ defmodule RemoteChain.NodeProxyTest do
 
       assert new_state.fallback == nil
       assert new_state.fallback_url == nil
+      refute Map.has_key?(new_state.lastblocks, @fallback_url)
     end
 
     test "keeps a fallback that has produced a block within the last 20 intervals" do
@@ -619,6 +620,33 @@ defmodule RemoteChain.NodeProxyTest do
       assert new_state.fallback == nil
       assert new_state.fallback_url == nil
     end
+
+    test "does not evict a fresh connection after a stale connection is replaced" do
+      # A stale timestamp is keyed by URL, not by WSConn pid. If removal keeps
+      # that timestamp, the next connection for the same URL is evicted at once.
+      url = "wss://base.example/ws"
+      stale_at = DateTime.add(DateTime.utc_now(), -1_800, :second)
+      old_conn = stub_data_connection(stale_at)
+
+      state = %NodeProxy{
+        chain: Chains.Base,
+        connections: %{url => old_conn},
+        fallback: nil,
+        fallback_url: nil,
+        lastblocks: %{url => {10, stale_at}}
+      }
+
+      state = NodeProxy.prune_stale_connections(state)
+
+      assert state.connections == %{}
+      refute Map.has_key?(state.lastblocks, url)
+
+      new_conn = stub_data_connection(DateTime.utc_now())
+      state = %{state | connections: %{url => new_conn}}
+
+      new_state = NodeProxy.prune_stale_connections(state)
+      assert new_state.connections == %{url => new_conn}
+    end
   end
 
   describe "prune_unresponsive_connections/1 watchdog" do
@@ -758,18 +786,29 @@ defmodule RemoteChain.NodeProxyTest do
     end
   end
 
+  defp stub_data_connection(lastblock_at) do
+    started_at = DateTime.utc_now()
+    {:ok, pid} = WSConnStateStub.start_state(%WSConn{started_at: started_at})
+    Globals.put({WSConn, pid}, :fake_conn)
+
+    :sys.replace_state(pid, fn %WSConn{} = wsconn ->
+      %{wsconn | lastblock_at: lastblock_at}
+    end)
+
+    on_exit(fn ->
+      Globals.pop({WSConn, pid})
+      if Process.alive?(pid), do: WSConn.close(pid)
+    end)
+
+    pid
+  end
+
   # Shared helpers for the `prune_stale_connections/1` eviction tests.
   # Mark the stub as ready (so the handshake-stale check is not the one
   # evicting) and use a started_at that is within the handshake timeout.
   # Only the data-staleness check should apply.
   defp with_fallback(state, lastblock_at) do
-    started_at = DateTime.utc_now()
-    {:ok, fallback} = WSConnStateStub.start_state(%WSConn{started_at: started_at})
-    Globals.put({WSConn, fallback}, :fake_conn)
-
-    :sys.replace_state(fallback, fn %WSConn{} = wsconn ->
-      %{wsconn | lastblock_at: lastblock_at}
-    end)
+    fallback = stub_data_connection(lastblock_at)
 
     %{state | fallback: fallback, fallback_url: "wss://fallback.example/oasis/mainnet/"}
   end