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

@@ -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

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.

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 },

vscode/src/client.ts

@@ -65,6 +65,7 @@ export class ImageClient {
   private revision = 0n;
   private closed = false;
   private ready = false;
+  private closing: Promise<void> = Promise.resolve();
 
   private constructor(
     private readonly socket: WebSocket,
@@ -313,15 +314,32 @@ export class ImageClient {
     }
   }
 
-  close(error?: Error): void {
+  close(error?: Error): Promise<void> {
-    if (this.closed) return;
+    if (this.closed) return this.closing;
     this.closed = true;
-    this.socket.terminate();
+    this.closing = new Promise<void>((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;
   }
 }

vscode/src/extension.ts

@@ -14,6 +14,12 @@ import {
 } from "./sourceProvider";
 import { SOURCE_ROOT } from "./sources";
 
+let closeConnections: (() => Promise<void>) | undefined;
+
+export async function deactivate(): Promise<void> {
+  await closeConnections?.();
+}
+
 class ActorTree implements vscode.TreeDataProvider<ActorSnapshot> {
   private readonly changes = new vscode.EventEmitter<void>();
   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<Promise<void>>();
+  const closeClient = (connection: ImageClient | undefined): Promise<void> => {
+    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();
+    const closed = disconnect();
-    endpoint = address;
     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?.();
       },
     },
   );

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<unknown[]> | 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) => {