diff --git a/packages/relay/src/cloudflare-adapter.test.ts b/packages/relay/src/cloudflare-adapter.test.ts index eeca70e700..e9aed32013 100644 --- a/packages/relay/src/cloudflare-adapter.test.ts +++ b/packages/relay/src/cloudflare-adapter.test.ts @@ -230,6 +230,39 @@ describe("RelayDurableObject control nudge/reset behavior", () => { }); }); +describe("RelayDurableObject data forwarding", () => { + it("forwards a client frame only to the newest daemon socket during replacement", () => { + const connectionId = "clt_replacing_daemon"; + const client = createMockSocket({ + version: "2", + role: "client", + connectionId, + serverId: "srv_test", + createdAt: 1, + }); + const daemonSockets = Array.from({ length: 6 }, (_, index) => + createMockSocket({ + version: "2", + role: "server", + connectionId, + serverId: "srv_test", + createdAt: index + 1, + }), + ); + const { state, setTagSockets } = createMockState(); + setTagSockets(`server:${connectionId}`, daemonSockets); + + const relay = new RelayDurableObject(state as unknown as DurableObjectStateArg); + relay.webSocketMessage(client as unknown as WebSocket, "create_agent_request"); + + for (const staleSocket of daemonSockets.slice(0, -1)) { + expect(staleSocket.send).not.toHaveBeenCalled(); + } + expect(daemonSockets.at(-1)?.send).toHaveBeenCalledOnce(); + expect(daemonSockets.at(-1)?.send).toHaveBeenCalledWith("create_agent_request"); + }); +}); + describe("relay worker endpoint routing", () => { it("routes missing v to legacy v1 isolated DO ids", async () => { const fetch = vi.fn( diff --git a/packages/relay/src/cloudflare-adapter.ts b/packages/relay/src/cloudflare-adapter.ts index e7e50c2309..4d77034649 100644 --- a/packages/relay/src/cloudflare-adapter.ts +++ b/packages/relay/src/cloudflare-adapter.ts @@ -86,6 +86,25 @@ function getString(record: Record, key: string): string | undef return typeof value === "string" ? value : undefined; } +function getSocketCreatedAt(ws: WebSocket): number { + const attachment = deserializeAttachment(ws); + if (!isRecord(attachment)) return 0; + return typeof attachment.createdAt === "number" ? attachment.createdAt : 0; +} + +function selectNewestSocket(sockets: WebSocket[]): WebSocket | null { + let newest: WebSocket | null = null; + let newestCreatedAt = Number.NEGATIVE_INFINITY; + for (const socket of sockets) { + const createdAt = getSocketCreatedAt(socket); + if (createdAt >= newestCreatedAt) { + newest = socket; + newestCreatedAt = createdAt; + } + } + return newest; +} + function getGlobalWebSocketPair(): (new () => WebSocketPair) | undefined { // Access WebSocketPair from global scope (Cloudflare Workers runtime) // Use Reflect to access global property without type assertions @@ -484,12 +503,11 @@ export class RelayDurableObject { this.bufferFrame(connectionId, message); return; } - for (const target of servers) { - try { - target.send(message); - } catch (error) { - console.error(`[Relay DO] Failed to forward client->server(${connectionId}):`, error); - } + const target = selectNewestSocket(servers); + try { + target?.send(message); + } catch (error) { + console.error(`[Relay DO] Failed to forward client->server(${connectionId}):`, error); } return; }