From ca0e14e5499e7800bba3b993011b6ca84f21bcfc Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Wed, 12 Aug 2026 10:04:16 +0000
Subject: [PATCH] fix(relay): order replacement sockets explicitly

Wall-clock timestamps can tie during rapid daemon socket replacement. Persist a per-connection generation so forwarding ownership follows acceptance order, including when replacing pre-generation sockets during rollout.
---
 packages/relay/src/cloudflare-adapter.test.ts | 37 ++++++++++++++++++-
 packages/relay/src/cloudflare-adapter.ts      | 20 +++++++---
 packages/relay/src/types.ts                   |  2 +
 3 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/packages/relay/src/cloudflare-adapter.test.ts b/packages/relay/src/cloudflare-adapter.test.ts
index e9aed32013..5bfb102bbb 100644
--- a/packages/relay/src/cloudflare-adapter.test.ts
+++ b/packages/relay/src/cloudflare-adapter.test.ts
@@ -246,7 +246,8 @@ describe("RelayDurableObject data forwarding", () => {
         role: "server",
         connectionId,
         serverId: "srv_test",
-        createdAt: index + 1,
+        generation: index + 1,
+        createdAt: 1,
       }),
     );
     const { state, setTagSockets } = createMockState();
@@ -261,6 +262,40 @@ describe("RelayDurableObject data forwarding", () => {
     expect(daemonSockets.at(-1)?.send).toHaveBeenCalledOnce();
     expect(daemonSockets.at(-1)?.send).toHaveBeenCalledWith("create_agent_request");
   });
+
+  it("prefers a generated replacement over rollout sockets without generations", () => {
+    const connectionId = "clt_rollout";
+    const client = createMockSocket({
+      version: "2",
+      role: "client",
+      connectionId,
+      serverId: "srv_test",
+      createdAt: 1,
+    });
+    const staleSocket = createMockSocket({
+      version: "2",
+      role: "server",
+      connectionId,
+      serverId: "srv_test",
+      createdAt: 2,
+    });
+    const replacementSocket = createMockSocket({
+      version: "2",
+      role: "server",
+      connectionId,
+      serverId: "srv_test",
+      generation: 1,
+      createdAt: 2,
+    });
+    const { state, setTagSockets } = createMockState();
+    setTagSockets(`server:${connectionId}`, [staleSocket, replacementSocket]);
+
+    const relay = new RelayDurableObject(state as unknown as DurableObjectStateArg);
+    relay.webSocketMessage(client as unknown as WebSocket, "create_agent_request");
+
+    expect(staleSocket.send).not.toHaveBeenCalled();
+    expect(replacementSocket.send).toHaveBeenCalledWith("create_agent_request");
+  });
 });
 
 describe("relay worker endpoint routing", () => {
diff --git a/packages/relay/src/cloudflare-adapter.ts b/packages/relay/src/cloudflare-adapter.ts
index 4d77034649..d65cb9dbc3 100644
--- a/packages/relay/src/cloudflare-adapter.ts
+++ b/packages/relay/src/cloudflare-adapter.ts
@@ -86,25 +86,29 @@ function getString(record: Record<string, unknown>, key: string): string | undef
   return typeof value === "string" ? value : undefined;
 }
 
-function getSocketCreatedAt(ws: WebSocket): number {
+function getSocketGeneration(ws: WebSocket): number {
   const attachment = deserializeAttachment(ws);
   if (!isRecord(attachment)) return 0;
-  return typeof attachment.createdAt === "number" ? attachment.createdAt : 0;
+  return typeof attachment.generation === "number" ? attachment.generation : 0;
 }
 
 function selectNewestSocket(sockets: WebSocket[]): WebSocket | null {
   let newest: WebSocket | null = null;
-  let newestCreatedAt = Number.NEGATIVE_INFINITY;
+  let newestGeneration = Number.NEGATIVE_INFINITY;
   for (const socket of sockets) {
-    const createdAt = getSocketCreatedAt(socket);
-    if (createdAt >= newestCreatedAt) {
+    const generation = getSocketGeneration(socket);
+    if (generation > newestGeneration) {
       newest = socket;
-      newestCreatedAt = createdAt;
+      newestGeneration = generation;
     }
   }
   return newest;
 }
 
+function nextSocketGeneration(sockets: WebSocket[]): number {
+  return sockets.reduce((highest, socket) => Math.max(highest, getSocketGeneration(socket)), 0) + 1;
+}
+
 function getGlobalWebSocketPair(): (new () => WebSocketPair) | undefined {
   // Access WebSocketPair from global scope (Cloudflare Workers runtime)
   // Use Reflect to access global property without type assertions
@@ -371,6 +375,9 @@ export class RelayDurableObject {
 
     const isServerControl = role === "server" && !resolvedConnectionId;
     const isServerData = role === "server" && !!resolvedConnectionId;
+    const serverDataGeneration = isServerData
+      ? nextSocketGeneration(this.state.getWebSockets(`server:${resolvedConnectionId}`))
+      : undefined;
 
     // Close any existing server-side connection with the same identity.
     // - server-control: single per serverId
@@ -396,6 +403,7 @@ export class RelayDurableObject {
       role,
       version: CURRENT_RELAY_VERSION,
       connectionId: resolvedConnectionId || null,
+      ...(serverDataGeneration === undefined ? {} : { generation: serverDataGeneration }),
       createdAt: Date.now(),
     };
     serializeAttachment(server, attachment);
diff --git a/packages/relay/src/types.ts b/packages/relay/src/types.ts
index 409435260a..9c8e46dd30 100644
--- a/packages/relay/src/types.ts
+++ b/packages/relay/src/types.ts
@@ -24,5 +24,7 @@ export interface RelaySessionAttachment {
    * independent socket + E2EE channel per connected connection.
    */
   connectionId?: string | null;
+  /** Monotonic acceptance order for v2 daemon data sockets sharing a connectionId. */
+  generation?: number;
   createdAt: number;
 }