From 89794b08dd4ba97616afc79169b6303078e6cd51 Mon Sep 17 00:00:00 2001 From: Steve Biedermann Date: Tue, 3 Nov 2026 17:24:09 +0100 Subject: [PATCH] fix: gracefully close extension WebSocket connections Use a code-1000 close handshake for intentional disconnects, wait for pending closes during deactivation, and bound shutdown to five seconds. Preserve immediate termination for failed connections and guard reconnects against stale sessions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 2 ++ vscode/README.md | 4 ++++ vscode/integration/server.test.ts | 13 +++++++++++++ vscode/src/client.ts | 24 +++++++++++++++++++++--- vscode/src/extension.ts | 31 ++++++++++++++++++++++++++----- vscode/test/client.test.ts | 20 ++++++++++++++++++++ 6 files changed, 86 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e7e3ae..82eb173 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,8 @@ All notable changes to this project will be documented in this file. ### Fixed +- Close extension connections with a WebSocket handshake on disconnect, + reconnect, and deactivation instead of terminating the transport immediately. - Exit the REPL on Ctrl+C and complete the WebSocket close handshake instead of returning to the prompt. - Complete the WebSocket close handshake when the CLI exits, preventing normal diff --git a/vscode/README.md b/vscode/README.md index d363643..aea4c3b 100644 --- a/vscode/README.md +++ b/vscode/README.md @@ -68,6 +68,10 @@ machine running the extension host, not necessarily the desktop. documents. Bursts are coalesced, and changes arriving during a fetch trigger a follow-up fetch. **Refresh Image** is also available on demand. - Open documents are retained and marked `connected: false` on disconnect. +- Explicit disconnect, reconnect, and extension deactivation complete the + WebSocket close handshake. Deactivation waits for outstanding closes; an + unresponsive peer is forcibly disconnected after five seconds. Protocol + corruption and request failures still terminate the connection without retry. Each connection gets distinct document URIs; reconnecting never replaces snapshots from an earlier connection or another image. Closing a document releases its cache; snapshots are not persisted across VS Code restarts. diff --git a/vscode/integration/server.test.ts b/vscode/integration/server.test.ts index bdc67c0..6a4d5ad 100644 --- a/vscode/integration/server.test.ts +++ b/vscode/integration/server.test.ts @@ -3,6 +3,19 @@ import test from "node:test"; import { ImageClient } from "../src/client"; import { control, spawnCounter, startImage } from "./server"; +test("extension disconnect and reconnect complete clean WebSocket handshakes", async (t) => { + const image = await startImage(); + t.after(() => image.stop()); + for (let attempt = 0; attempt < 3; attempt++) { + const client = await ImageClient.connect(image.address, (error) => { + assert.equal(error, undefined); + }); + await client.listActors(); + await client.close(); + } + assert.doesNotMatch(image.diagnostics(), /protocol error|reset without closing handshake/i); +}); + test( "TypeScript inspection shares the real Rust image across clients and reconnects", { timeout: 10000 }, diff --git a/vscode/src/client.ts b/vscode/src/client.ts index 9550295..821b95e 100644 --- a/vscode/src/client.ts +++ b/vscode/src/client.ts @@ -65,6 +65,7 @@ export class ImageClient { private revision = 0n; private closed = false; private ready = false; + private closing: Promise = Promise.resolve(); private constructor( private readonly socket: WebSocket, @@ -313,15 +314,32 @@ export class ImageClient { } } - close(error?: Error): void { - if (this.closed) return; + close(error?: Error): Promise { + if (this.closed) return this.closing; this.closed = true; - this.socket.terminate(); + this.closing = new Promise((resolve) => { + if (this.socket.readyState === WebSocket.CLOSED) { + resolve(); + return; + } + const timer = setTimeout(() => this.socket.terminate(), 5000); + timer.unref(); + this.socket.once("close", () => { + clearTimeout(timer); + resolve(); + }); + if (error || this.socket.readyState === WebSocket.CONNECTING) { + this.socket.terminate(); + } else { + this.socket.close(1000); + } + }); if (this.pending) { clearTimeout(this.pending.timer); this.pending.reject(error ?? new Error("Disconnected from image")); this.pending = undefined; } if (this.ready) this.onDisconnect(error); + return this.closing; } } diff --git a/vscode/src/extension.ts b/vscode/src/extension.ts index 57c1e73..5022f08 100644 --- a/vscode/src/extension.ts +++ b/vscode/src/extension.ts @@ -14,6 +14,12 @@ import { } from "./sourceProvider"; import { SOURCE_ROOT } from "./sources"; +let closeConnections: (() => Promise) | undefined; + +export async function deactivate(): Promise { + await closeConnections?.(); +} + class ActorTree implements vscode.TreeDataProvider { private readonly changes = new vscode.EventEmitter(); readonly onDidChangeTreeData = this.changes.event; @@ -140,6 +146,18 @@ export function activate(context: vscode.ExtensionContext): void { let endpoint = ""; let disposed = false; let automaticRefresh: RefreshQueue | undefined; + const closing = new Set>(); + const closeClient = (connection: ImageClient | undefined): Promise => { + if (!connection) return Promise.resolve(); + const closed = connection.close(); + closing.add(closed); + void closed.then(() => closing.delete(closed)); + return closed; + }; + closeConnections = async () => { + await closeClient(client); + await Promise.all(closing); + }; const report = (error: unknown) => { const failure = error instanceof Error ? error : new Error(String(error)); @@ -168,9 +186,10 @@ export function activate(context: vscode.ExtensionContext): void { documents.disconnect(session); sources.store.detach(); session++; - previous?.close(); + const closed = closeClient(previous); tree.replace([]); connectionStatus(false); + return closed; }; const refresh = async () => { @@ -245,9 +264,11 @@ export function activate(context: vscode.ExtensionContext): void { })); if (address === undefined || disposed) return; parseEndpoint(address); - disconnect(); - endpoint = address; + const closed = disconnect(); const currentSession = session; + await closed; + if (currentSession !== session || disposed) return; + endpoint = address; view.message = `Connecting to ${address}...`; let connected: ImageClient; const updates = new RefreshQueue(async () => { @@ -276,7 +297,7 @@ export function activate(context: vscode.ExtensionContext): void { throw error; } if (currentSession !== session || disposed) { - connected.close(); + await closeClient(connected); return; } client = connected; @@ -662,7 +683,7 @@ export function activate(context: vscode.ExtensionContext): void { dispose: () => { disposed = true; automaticRefresh?.dispose(); - client?.close(); + void closeConnections?.(); }, }, ); diff --git a/vscode/test/client.test.ts b/vscode/test/client.test.ts index 21b7053..5108c60 100644 --- a/vscode/test/client.test.ts +++ b/vscode/test/client.test.ts @@ -25,6 +25,26 @@ test("only numeric loopback endpoints and valid ports are accepted", () => { } }); +test("normal disconnect completes a code-1000 close handshake and rejects further work", async (t) => { + let closed: Promise | undefined; + const address = await fakeImage(t, (request, socket) => { + closed = once(socket, "close"); + socket.send(response(request, { kind: "actors", data: [] })); + }); + let disconnects = 0; + const client = await ImageClient.connect(address, (error) => { + assert.equal(error, undefined); + disconnects++; + }); + await client.listActors(); + await client.close(); + assert.ok(closed); + assert.equal((await closed)[0], 1000); + await client.close(); + assert.equal(disconnects, 1); + await assert.rejects(client.listActors(), /Not connected/); +}); + test("fragmented UTF-8 replies and simultaneous reads preserve correlation and values", async (t) => { const operations: string[] = []; const address = await fakeImage(t, (request, socket) => {