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) => {