From 1963f485202e82368d96cf7a4b3964b27bede1f9 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Mon, 17 Aug 2026 19:26:09 -0700
Subject: [PATCH] fix(gateway): materialize the WaveKV winner in proxy state

---
 dstack/gateway/src/main_service.rs       | 28 +++++++++------
 dstack/gateway/src/main_service/tests.rs | 45 ++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs
index 6b9763c23..825041fd4 100644
--- a/dstack/gateway/src/main_service.rs
+++ b/dstack/gateway/src/main_service.rs
@@ -1100,7 +1100,7 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()>
     }
 
     for (instance_id, data) in instances {
-        let new_info = InstanceInfo {
+        let mut new_info = InstanceInfo {
             id: instance_id.clone(),
             app_id: data.app_id.clone(),
             ip: data.ip,
@@ -1114,24 +1114,32 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()>
             connections: Default::default(),
         };
 
-        let old_ip = state.state.instances.get(&instance_id).map(|e| e.ip);
-        if let Some(existing) = state.state.instances.get(&instance_id) {
+        let existing = state.state.instances.get(&instance_id).cloned();
+        if let Some(existing) = &existing {
             // Check if wg config needs update
             if existing.public_key != data.public_key || existing.ip != data.ip {
                 wg_changed = true;
             }
-            // Only update if remote is newer (based on reg_time)
-            if data.reg_time <= encode_ts(existing.reg_time) {
-                continue;
-            }
+            // WaveKV has already selected the winning value. Materialize it
+            // unconditionally instead of applying another LWW rule here.
+            new_info.connections = existing.connections.clone();
         } else {
             wg_changed = true;
         }
 
         // Release old IP if it changed (prevent IP leak)
-        if let Some(old_ip) = old_ip {
-            if old_ip != data.ip {
-                state.state.allocated_addresses.remove(&old_ip);
+        if let Some(existing) = &existing {
+            if existing.ip != data.ip {
+                state.state.allocated_addresses.remove(&existing.ip);
+            }
+            if existing.app_id != data.app_id {
+                if let Some(app_instances) = state.state.apps.get_mut(&existing.app_id) {
+                    app_instances.remove(&instance_id);
+                    if app_instances.is_empty() {
+                        state.state.apps.remove(&existing.app_id);
+                        state.state.top_n.remove(&existing.app_id);
+                    }
+                }
             }
         }
         state.state.allocated_addresses.insert(data.ip);
diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs
index 7135c9007..4280da273 100644
--- a/dstack/gateway/src/main_service/tests.rs
+++ b/dstack/gateway/src/main_service/tests.rs
@@ -7,6 +7,7 @@ use crate::config::{load_config_figment, Config, MutualConfig};
 use crate::kv::PortFlags;
 use crate::proxy::port_policy::is_port_allowed;
 use base64::Engine as _;
+use std::sync::atomic::Ordering;
 use tempfile::TempDir;
 
 struct TestState {
@@ -557,6 +558,50 @@ async fn an_instance_deleted_on_another_node_stops_being_routable_here() {
         .contains(&"10.0.0.40".parse().unwrap()));
 }
 
+#[tokio::test]
+async fn proxy_state_adopts_the_wavekv_winner_regardless_of_value_reg_time() {
+    let state = create_test_state().await;
+    sync_from_peer_at(
+        &state,
+        "contended",
+        "10.0.0.40",
+        &test_pubkey("old-key"),
+        300,
+    );
+    reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap();
+    state
+        .lock()
+        .state
+        .instances
+        .get("contended")
+        .unwrap()
+        .connections
+        .store(7, Ordering::Relaxed);
+
+    // This is the value WaveKV selected using its own entry metadata. Its
+    // payload timestamp is older, so comparing reg_time again would leave the
+    // data plane permanently materializing a losing value.
+    sync_from_peer_at(
+        &state,
+        "contended",
+        "10.0.0.41",
+        &test_pubkey("winner-key"),
+        100,
+    );
+    reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap();
+
+    let proxy = state.lock();
+    let instance = &proxy.state.instances["contended"];
+    assert_eq!(instance.ip, "10.0.0.41".parse::<Ipv4Addr>().unwrap());
+    assert_eq!(instance.public_key, test_pubkey("winner-key"));
+    assert_eq!(encode_ts(instance.reg_time), 100);
+    assert_eq!(instance.num_connections(), 7);
+    assert!(!proxy
+        .state
+        .allocated_addresses
+        .contains(&"10.0.0.40".parse().unwrap()));
+}
+
 #[tokio::test]
 async fn a_local_registration_survives_a_reload_that_cannot_see_it_yet() {
     let state = create_test_state().await;