Skip to content

Commit 69df291

Browse files
EhabYEhabY
committed
Only trigger reconnection if the WS is not connected
1 parent 2e27a6b commit 69df291

File tree

4 files changed

+113
-27
lines changed

4 files changed

+113
-27
lines changed

‎src/api/coderApi.ts‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ import {
4646
type OneWayWebSocketInit,
4747
} from "../websocket/oneWayWebSocket";
4848
import {
49+
ConnectionState,
4950
ReconnectingWebSocket,
5051
type SocketFactory,
5152
} from "../websocket/reconnectingWebSocket";
@@ -164,21 +165,23 @@ export class CoderApi extends Api implements vscode.Disposable {
164165

165166
/**
166167
* Watch for configuration changes that affect WebSocket connections.
167-
* When any watched setting changes, all active WebSockets are reconnected.
168+
* Only reconnects DISCONNECTED sockets since they require an explicit reconnect() call.
169+
* Other states will pick up settings naturally.
168170
*/
169171
private watchConfigChanges(): vscode.Disposable {
170172
const settings = webSocketConfigSettings.map((setting) => ({
171173
setting,
172174
getValue: () => vscode.workspace.getConfiguration().get(setting),
173175
}));
174176
return watchConfigurationChanges(settings, () => {
175-
if (this.reconnectingSockets.size > 0) {
176-
this.output.info(
177-
`Configuration changed, reconnecting ${this.reconnectingSockets.size} WebSocket(s)`,
178-
);
179-
for (const socket of this.reconnectingSockets) {
180-
socket.reconnect();
181-
}
177+
const socketsToReconnect = [...this.reconnectingSockets].filter(
178+
(socket) => socket.state === ConnectionState.DISCONNECTED,
179+
);
180+
this.output.debug(
181+
`Configuration changed, ${socketsToReconnect.length}/${this.reconnectingSockets.size} socket(s) in DISCONNECTED state`,
182+
);
183+
for (const socket of socketsToReconnect) {
184+
socket.reconnect();
182185
}
183186
});
184187
}

‎src/websocket/reconnectingWebSocket.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ export class ReconnectingWebSocket<
197197
/**
198198
* Returns the current connection state.
199199
*/
200-
get state(): string {
200+
get state(): ConnectionState {
201201
return this.#state;
202202
}
203203

‎test/unit/api/coderApi.test.ts‎

Lines changed: 52 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -591,9 +591,11 @@ describe("CoderApi", () => {
591591

592592
const setupAutoOpeningWebSocket = () => {
593593
const sockets: Array<Partial<Ws>> = [];
594+
const handlers: Record<string, (...args: unknown[]) => void> = {};
594595
vi.mocked(Ws).mockImplementation(function (url: string | URL) {
595596
const mockWs = createMockWebSocket(String(url), {
596-
on: vi.fn((event, handler) => {
597+
on: vi.fn((event: string, handler: (...args: unknown[]) => void) => {
598+
handlers[event] = handler;
597599
if (event === "open") {
598600
setImmediate(() => handler());
599601
}
@@ -603,12 +605,12 @@ describe("CoderApi", () => {
603605
sockets.push(mockWs);
604606
return mockWs as Ws;
605607
});
606-
return sockets;
608+
return { sockets, handlers };
607609
};
608610

609611
describe("Reconnection on Host/Token Changes", () => {
610612
it("triggers reconnection when session token changes", async () => {
611-
const sockets = setupAutoOpeningWebSocket();
613+
const { sockets } = setupAutoOpeningWebSocket();
612614
api = createApi(CODER_URL, AXIOS_TOKEN);
613615
await api.watchAgentMetadata(AGENT_ID);
614616

@@ -623,7 +625,7 @@ describe("CoderApi", () => {
623625
});
624626

625627
it("triggers reconnection when host changes", async () => {
626-
const sockets = setupAutoOpeningWebSocket();
628+
const { sockets } = setupAutoOpeningWebSocket();
627629
api = createApi(CODER_URL, AXIOS_TOKEN);
628630
const wsWrap = await api.watchAgentMetadata(AGENT_ID);
629631
expect(wsWrap.url).toContain(CODER_URL.replace("http", "ws"));
@@ -642,7 +644,7 @@ describe("CoderApi", () => {
642644
});
643645

644646
it("does not reconnect when token or host are unchanged", async () => {
645-
const sockets = setupAutoOpeningWebSocket();
647+
const { sockets } = setupAutoOpeningWebSocket();
646648
api = createApi(CODER_URL, AXIOS_TOKEN);
647649
await api.watchAgentMetadata(AGENT_ID);
648650

@@ -655,7 +657,7 @@ describe("CoderApi", () => {
655657
});
656658

657659
it("suspends sockets when host is set to empty string (logout)", async () => {
658-
const sockets = setupAutoOpeningWebSocket();
660+
const { sockets } = setupAutoOpeningWebSocket();
659661
api = createApi(CODER_URL, AXIOS_TOKEN);
660662
await api.watchAgentMetadata(AGENT_ID);
661663

@@ -668,7 +670,7 @@ describe("CoderApi", () => {
668670
});
669671

670672
it("does not reconnect when setting token after clearing host", async () => {
671-
const sockets = setupAutoOpeningWebSocket();
673+
const { sockets } = setupAutoOpeningWebSocket();
672674
api = createApi(CODER_URL, AXIOS_TOKEN);
673675
await api.watchAgentMetadata(AGENT_ID);
674676

@@ -682,7 +684,7 @@ describe("CoderApi", () => {
682684
});
683685

684686
it("setCredentials sets both host and token together", async () => {
685-
const sockets = setupAutoOpeningWebSocket();
687+
const { sockets } = setupAutoOpeningWebSocket();
686688
api = createApi(CODER_URL, AXIOS_TOKEN);
687689
await api.watchAgentMetadata(AGENT_ID);
688690

@@ -695,7 +697,7 @@ describe("CoderApi", () => {
695697
});
696698

697699
it("setCredentials suspends when host is cleared", async () => {
698-
const sockets = setupAutoOpeningWebSocket();
700+
const { sockets } = setupAutoOpeningWebSocket();
699701
api = createApi(CODER_URL, AXIOS_TOKEN);
700702
await api.watchAgentMetadata(AGENT_ID);
701703

@@ -796,29 +798,26 @@ describe("CoderApi", () => {
796798
describe("Configuration Change Reconnection", () => {
797799
const tick = () => new Promise((resolve) => setImmediate(resolve));
798800

799-
it("reconnects sockets when watched config value changes", async () => {
801+
it("does not reconnect connected sockets when config value changes", async () => {
800802
mockConfig.set("coder.insecure", false);
801-
const sockets = setupAutoOpeningWebSocket();
803+
const { sockets } = setupAutoOpeningWebSocket();
802804
api = createApi(CODER_URL, AXIOS_TOKEN);
803805
await api.watchAgentMetadata(AGENT_ID);
804806
await tick(); // Wait for open event to fire (socket becomes CONNECTED)
805807

806808
mockConfig.set("coder.insecure", true);
807809
await tick();
808810

809-
expect(sockets[0].close).toHaveBeenCalledWith(
810-
1000,
811-
"Replacing connection",
812-
);
813-
expect(sockets).toHaveLength(2);
811+
expect(sockets[0].close).not.toHaveBeenCalled();
812+
expect(sockets).toHaveLength(1);
814813
});
815814

816815
it.each([
817816
["unchanged value", "coder.insecure", false],
818817
["unrelated setting", "unrelated.setting", "new-value"],
819818
])("does not reconnect for %s", async (_desc, key, value) => {
820819
mockConfig.set("coder.insecure", false);
821-
const sockets = setupAutoOpeningWebSocket();
820+
const { sockets } = setupAutoOpeningWebSocket();
822821
api = createApi(CODER_URL, AXIOS_TOKEN);
823822
await api.watchAgentMetadata(AGENT_ID);
824823
await tick(); // Wait for open event to fire (socket becomes CONNECTED)
@@ -832,7 +831,7 @@ describe("CoderApi", () => {
832831

833832
it("stops watching after dispose", async () => {
834833
mockConfig.set("coder.insecure", false);
835-
const sockets = setupAutoOpeningWebSocket();
834+
const { sockets } = setupAutoOpeningWebSocket();
836835
api = createApi(CODER_URL, AXIOS_TOKEN);
837836
await api.watchAgentMetadata(AGENT_ID);
838837
await tick(); // Wait for open event to fire (socket becomes CONNECTED)
@@ -843,6 +842,41 @@ describe("CoderApi", () => {
843842

844843
expect(sockets).toHaveLength(1);
845844
});
845+
846+
it("does not reconnect sockets in AWAITING_RETRY state when config changes", async () => {
847+
mockConfig.set("coder.insecure", false);
848+
const { sockets, handlers } = setupAutoOpeningWebSocket();
849+
api = createApi(CODER_URL, AXIOS_TOKEN);
850+
await api.watchAgentMetadata(AGENT_ID);
851+
852+
// Trigger close with abnormal code to put socket in AWAITING_RETRY
853+
handlers["close"]?.({ code: 1006, reason: "Abnormal closure" });
854+
await tick();
855+
856+
mockConfig.set("coder.insecure", true);
857+
await tick();
858+
859+
// AWAITING_RETRY will naturally retry, so no config-triggered reconnect needed
860+
expect(sockets).toHaveLength(1);
861+
});
862+
863+
it("reconnects sockets in DISCONNECTED state when config changes", async () => {
864+
mockConfig.set("coder.insecure", false);
865+
const { sockets, handlers } = setupAutoOpeningWebSocket();
866+
api = createApi(CODER_URL, AXIOS_TOKEN);
867+
await api.watchAgentMetadata(AGENT_ID);
868+
await tick();
869+
870+
// Trigger close with unrecoverable code to put socket in DISCONNECTED
871+
handlers["close"]?.({ code: 1002, reason: "Protocol error" });
872+
await tick();
873+
874+
mockConfig.set("coder.insecure", true);
875+
await tick();
876+
877+
// Only DISCONNECTED sockets get reconnected by config changes
878+
expect(sockets).toHaveLength(2);
879+
});
846880
});
847881
});
848882

‎test/unit/websocket/reconnectingWebSocket.test.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,55 @@ describe("ReconnectingWebSocket", () => {
510510
});
511511

512512
describe("Error Handling", () => {
513+
it("error event when CONNECTED schedules retry", async () => {
514+
const { ws, sockets } = await createReconnectingWebSocket();
515+
516+
sockets[0].fireOpen();
517+
expect(ws.state).toBe(ConnectionState.CONNECTED);
518+
519+
sockets[0].fireError(new Error("Connection lost"));
520+
expect(ws.state).toBe(ConnectionState.AWAITING_RETRY);
521+
522+
await vi.advanceTimersByTimeAsync(300);
523+
expect(sockets).toHaveLength(2);
524+
525+
ws.close();
526+
});
527+
528+
it("error event when DISCONNECTED is ignored", async () => {
529+
const { ws, sockets } = await createReconnectingWebSocket();
530+
531+
sockets[0].fireOpen();
532+
expect(ws.state).toBe(ConnectionState.CONNECTED);
533+
534+
ws.disconnect();
535+
expect(ws.state).toBe(ConnectionState.DISCONNECTED);
536+
537+
// Error after disconnect should be ignored
538+
sockets[0].fireError(new Error("Connection lost"));
539+
expect(ws.state).toBe(ConnectionState.DISCONNECTED);
540+
541+
// No reconnection should be scheduled
542+
await vi.advanceTimersByTimeAsync(10000);
543+
expect(sockets).toHaveLength(1);
544+
545+
ws.close();
546+
});
547+
548+
it("multiple errors while AWAITING_RETRY only creates one reconnection", async () => {
549+
const { ws, sockets } = await createReconnectingWebSocket();
550+
sockets[0].fireOpen();
551+
552+
sockets[0].fireError(new Error("First error"));
553+
sockets[0].fireError(new Error("Second error"));
554+
sockets[0].fireError(new Error("Third error"));
555+
556+
await vi.advanceTimersByTimeAsync(300);
557+
expect(sockets).toHaveLength(2);
558+
559+
ws.close();
560+
});
561+
513562
it("schedules retry when socket factory throws error", async () => {
514563
const sockets: MockSocket[] = [];
515564
let shouldFail = false;

0 commit comments

Comments
 (0)