@ekodb/ekodb-client 0.19.0 → 0.20.0

package/dist/client.js

@@ -69,6 +69,19 @@ var MergeStrategy;
     MergeStrategy["LatestOnly"] = "LatestOnly";
     MergeStrategy["Interleaved"] = "Interleaved";
 })(MergeStrategy || (exports.MergeStrategy = MergeStrategy = {}));
+/**
+ * Strip trailing slashes from a base URL so path concatenation
+ * (`${base}/api/...`) never yields a double-slash path. Uses a linear scan
+ * rather than a regex like `/\/+$/`, which CodeQL flags as polynomial-time
+ * backtracking on caller-supplied input.
+ */
+function stripTrailingSlashes(url) {
+    let end = url.length;
+    while (end > 0 && url.charCodeAt(end - 1) === 47 /* "/" */) {
+        end--;
+    }
+    return end === url.length ? url : url.slice(0, end);
+}
 class EkoDBClient {
     constructor(config, apiKey) {
         this.token = null;
@@ -76,14 +89,16 @@ class EkoDBClient {
         this.rateLimitInfo = null;
         // Support both old (baseURL, apiKey) and new (config object) signatures
         if (typeof config === "string") {
-            this.baseURL = config;
+            // Strip trailing slashes so `${baseURL}/api/...` never produces a
+            // double-slash path (some servers/proxies reject `//api/...`).
+            this.baseURL = stripTrailingSlashes(config);
             this.apiKey = apiKey;
             this.shouldRetry = true;
             this.maxRetries = 3;
             this.format = SerializationFormat.MessagePack; // Default to MessagePack for 2-3x performance
         }
         else {
-            this.baseURL = config.baseURL;
+            this.baseURL = stripTrailingSlashes(config.baseURL);
             this.apiKey = config.apiKey;
             this.shouldRetry = config.shouldRetry ?? true;
             this.maxRetries = config.maxRetries ?? 3;
@@ -223,6 +238,38 @@ class EkoDBClient {
     sleep(seconds) {
         return new Promise((resolve) => setTimeout(resolve, seconds * 1000));
     }
+    /**
+     * Parse a `Retry-After` header into a non-negative delay in seconds.
+     *
+     * Per RFC 9110 the value is either delay-seconds (an integer) or an
+     * HTTP-date. Anything that doesn't resolve to a finite, non-negative number
+     * (missing header, garbage, a past date) falls back to `defaultSecs`.
+     */
+    parseRetryAfter(header, defaultSecs = 60) {
+        if (!header)
+            return defaultSecs;
+        // delay-seconds form: a bare integer.
+        const secs = Number(header.trim());
+        if (Number.isFinite(secs))
+            return Math.max(0, secs);
+        // HTTP-date form: compute the delay from now.
+        const dateMs = Date.parse(header);
+        if (Number.isFinite(dateMs)) {
+            return Math.max(0, (dateMs - Date.now()) / 1000);
+        }
+        return defaultSecs;
+    }
+    /**
+     * Backoff delay (in seconds) for a 0-indexed retry attempt: a capped
+     * exponential schedule (0.2s → 5s) with full jitter, so concurrent clients
+     * don't retry in lockstep. Returns a value in [d/2, d].
+     */
+    backoffSeconds(attempt) {
+        const base = 0.2;
+        const max = 5;
+        const d = Math.min(base * Math.pow(2, Math.max(0, attempt)), max);
+        return d / 2 + Math.random() * (d / 2);
+    }
     /**
      * Helper to determine if a path should use JSON
      * Only CRUD operations (insert/update/delete/batch) use MessagePack
@@ -292,10 +339,13 @@ class EkoDBClient {
             }
             // Handle rate limiting (429)
             if (response.status === 429) {
-                const retryAfter = parseInt(response.headers.get("retry-after") || "60", 10);
+                const retryAfter = this.parseRetryAfter(response.headers.get("retry-after"));
                 if (this.shouldRetry && attempt < this.maxRetries) {
-                    console.log(`Rate limited. Retrying after ${retryAfter} seconds...`);
-                    await this.sleep(retryAfter);
+                    // Honor the server's Retry-After, but cap it so a hostile/large value
+                    // can't pin the client for minutes.
+                    const wait = Math.min(retryAfter, 60);
+                    console.log(`Rate limited. Retrying after ${wait} seconds...`);
+                    await this.sleep(wait);
                     return this.makeRequest(method, path, data, attempt + 1, forceJson);
                 }
                 throw new RateLimitError(retryAfter);
@@ -311,8 +361,8 @@ class EkoDBClient {
             if (response.status === 503 &&
                 this.shouldRetry &&
                 attempt < this.maxRetries) {
-                const retryDelay = 10;
-                console.log(`Service unavailable. Retrying after ${retryDelay} seconds...`);
+                const retryDelay = this.backoffSeconds(attempt);
+                console.log(`Service unavailable. Retrying after ${retryDelay.toFixed(2)}s...`);
                 await this.sleep(retryDelay);
                 return this.makeRequest(method, path, data, attempt + 1, forceJson);
             }
@@ -325,8 +375,8 @@ class EkoDBClient {
             if (error instanceof TypeError &&
                 this.shouldRetry &&
                 attempt < this.maxRetries) {
-                const retryDelay = 3;
-                console.log(`Network error. Retrying after ${retryDelay} seconds...`);
+                const retryDelay = this.backoffSeconds(attempt);
+                console.log(`Network error. Retrying after ${retryDelay.toFixed(2)}s...`);
                 await this.sleep(retryDelay);
                 return this.makeRequest(method, path, data, attempt + 1, forceJson);
             }
@@ -1125,10 +1175,10 @@ class EkoDBClient {
         const stream = new EventStream();
         (async () => {
             try {
-                let token = this.getToken();
+                let token = await this.getToken();
                 if (!token) {
                     await this.refreshToken();
-                    token = this.getToken();
+                    token = await this.getToken();
                 }
                 const url = `${this.baseURL}/api/chat/${chatId}/messages/stream`;
                 const response = await fetch(url, {
@@ -1149,13 +1199,12 @@ class EkoDBClient {
                     stream.close();
                     return;
                 }
-                const body = await response.text();
-                for (const line of body.split("\n")) {
+                const emitLine = (line) => {
                     if (!line.startsWith("data:"))
-                        continue;
+                        return;
                     const dataStr = line.slice(5).trim();
                     if (!dataStr)
-                        continue;
+                        return;
                     try {
                         const eventData = JSON.parse(dataStr);
                         if (eventData.error) {
@@ -1184,6 +1233,33 @@ class EkoDBClient {
                     catch {
                         // skip malformed SSE data
                     }
+                };
+                const reader = response.body?.getReader?.();
+                if (reader) {
+                    // True incremental streaming: decode and emit each SSE line as soon as
+                    // it arrives, rather than buffering the entire response body first.
+                    const decoder = new TextDecoder();
+                    let buffer = "";
+                    for (;;) {
+                        const { done, value } = await reader.read();
+                        if (done)
+                            break;
+                        buffer += decoder.decode(value, { stream: true });
+                        let nl;
+                        while ((nl = buffer.indexOf("\n")) >= 0) {
+                            emitLine(buffer.slice(0, nl));
+                            buffer = buffer.slice(nl + 1);
+                        }
+                    }
+                    buffer += decoder.decode();
+                    if (buffer)
+                        emitLine(buffer);
+                }
+                else {
+                    // Fallback for environments/tests without a readable body stream.
+                    const body = await response.text();
+                    for (const line of body.split("\n"))
+                        emitLine(line);
                 }
                 stream.close();
             }
@@ -1736,10 +1812,18 @@ class EkoDBClient {
         return stream;
     }
     /**
-     * Create a WebSocket client
+     * Create a WebSocket client.
+     *
+     * The token is supplied as a provider bound to this client's
+     * {@link getToken}, so every (re)connect re-evaluates (and proactively
+     * refreshes) the auth token instead of snapshotting it once. This means a
+     * reconnect after a token rotation uses the current token.
+     *
+     * @param wsURL - The WebSocket URL (e.g. `wss://host`); `/api/ws` is appended if absent.
+     * @param options - Optional reconnect/timeout tunables.
      */
-    websocket(wsURL) {
-        return new WebSocketClient(wsURL, this.token);
+    websocket(wsURL, options) {
+        return new WebSocketClient(wsURL, () => this.getToken(), options);
     }
     // ========== RAG Helper Methods ==========
     /**
@@ -1988,50 +2072,108 @@ function extractRecordId(record, extraCandidates = []) {
         const val = record[key];
         if (typeof val === "string")
             return val;
-        if (val && typeof val === "object" && "value" in val)
+        // Unwrap only a genuine typed wrapper (both "type" and "value"), matching
+        // getValue's rule so a user object like { value: 1 } isn't mistaken for one.
+        if (val && typeof val === "object" && "type" in val && "value" in val)
             return String(val.value);
     }
     for (const key of ["id", "_id"]) {
         const val = record[key];
         if (typeof val === "string")
             return val;
-        if (val && typeof val === "object" && "value" in val)
+        // Unwrap only a genuine typed wrapper (both "type" and "value"), matching
+        // getValue's rule so a user object like { value: 1 } isn't mistaken for one.
+        if (val && typeof val === "object" && "type" in val && "value" in val)
             return String(val.value);
     }
     return undefined;
 }
 class WebSocketClient {
-    constructor(wsURL, token) {
+    /**
+     * @param wsURL - WebSocket URL; `/api/ws` is appended if absent.
+     * @param token - A static token string OR a {@link TokenProvider} function
+     *   re-evaluated on every (re)connect (so a refreshed token is used after a drop).
+     * @param options - Optional reconnect/timeout tunables.
+     */
+    constructor(wsURL, token, options = {}) {
         this.ws = null;
         this.dispatcherRunning = false;
         this.schemaCache = null;
+        // Reconnect state
+        /** Set while close() is in progress so the close handler doesn't reconnect. */
+        this.closed = false;
+        this.reconnectAttempts = 0;
+        this.reconnecting = false;
+        this.connectPromise = null;
         // Dispatcher state
         this.pendingRequests = new Map();
         this.subscriptions = new Map();
+        /** Bookkeeping so subscriptions can be replayed on reconnect. */
+        this.subscriptionParams = new Map();
         this.chatStreams = new Map();
         this.registerToolsAck = null;
         this.messageCounter = 0;
-        this.wsURL = wsURL;
-        this.token = token;
+        // Strip trailing slashes so appending `/api/ws` can't yield `//api/ws`,
+        // which warp's exact path match (`api / ws`) would reject.
+        this.wsURL = stripTrailingSlashes(wsURL);
+        this.tokenProvider = typeof token === "function" ? token : () => token;
+        this.autoReconnect = options.autoReconnect ?? true;
+        this.reconnectInitialDelayMs = options.reconnectInitialDelayMs ?? 200;
+        this.reconnectMaxDelayMs = options.reconnectMaxDelayMs ?? 5000;
+        this.reconnectMaxAttempts = options.reconnectMaxAttempts ?? 0;
+        this.requestTimeoutMs = options.requestTimeoutMs ?? 30000;
     }
     genMessageId() {
         const counter = this.messageCounter++;
         return `${Date.now()}-${counter}-${Math.random().toString(36).slice(2, 8)}`;
     }
     /**
-     * Connect and start the dispatcher.
+     * Compute the capped exponential backoff (with jitter) for a reconnect
+     * attempt. attempt 0 -> ~initial, growing x2 each time up to the max cap.
+     * Jitter is +/-25% to avoid thundering-herd reconnect storms.
+     * @internal exposed for testing
+     */
+    computeBackoff(attempt) {
+        const base = Math.min(this.reconnectInitialDelayMs * 2 ** attempt, this.reconnectMaxDelayMs);
+        const jitter = base * 0.25 * (Math.random() * 2 - 1);
+        return Math.max(0, Math.round(base + jitter));
+    }
+    /**
+     * Connect and start the dispatcher. Re-evaluates the token provider so the
+     * current/refreshed token is used for this socket.
      */
     async ensureConnected() {
         if (this.ws && this.dispatcherRunning)
             return;
+        // Coalesce concurrent connect attempts onto a single in-flight promise.
+        if (this.connectPromise)
+            return this.connectPromise;
+        // Clear the intentional-close flag only for user-initiated connects. During
+        // a reconnect cycle this stays untouched so a concurrent close() can't be
+        // undone and have the reconnect proceed against the user's intent.
+        if (!this.reconnecting)
+            this.closed = false;
+        this.connectPromise = this.openSocket().finally(() => {
+            this.connectPromise = null;
+        });
+        return this.connectPromise;
+    }
+    async openSocket() {
         const WebSocket = (await Promise.resolve().then(() => __importStar(require("ws")))).default;
         let url = this.wsURL;
         if (!url.endsWith("/api/ws")) {
             url += "/api/ws";
         }
+        // Re-evaluate the token on every (re)connect — never a stale snapshot.
+        const token = await this.tokenProvider();
+        if (!token) {
+            // Fail fast with a clear error instead of sending `Bearer null`, which
+            // would surface as a confusing 401 from the server.
+            throw new Error("WebSocket auth token is unavailable (the token provider returned null/empty)");
+        }
         this.ws = new WebSocket(url, {
             headers: {
-                Authorization: `Bearer ${this.token}`,
+                Authorization: `Bearer ${token}`,
             },
         });
         await new Promise((resolve, reject) => {
@@ -2044,7 +2186,13 @@ class WebSocketClient {
         if (this.dispatcherRunning)
             return;
         this.dispatcherRunning = true;
-        this.ws.on("message", (data) => {
+        // Capture the socket this dispatcher is bound to. After a reconnect, the old
+        // socket may still emit late close/error events; ignore them so they don't
+        // tear down the replacement connection.
+        const socket = this.ws;
+        socket.on("message", (data) => {
+            if (this.ws !== socket)
+                return;
             try {
                 const msg = JSON.parse(data.toString());
                 this.routeMessage(msg);
@@ -2053,26 +2201,149 @@ class WebSocketClient {
                 // Ignore malformed messages
             }
         });
-        this.ws.on("close", () => {
-            this.dispatcherRunning = false;
-            // Notify all pending requests
-            for (const [, pending] of this.pendingRequests) {
-                pending.reject(new Error("WebSocket connection closed"));
-            }
-            this.pendingRequests.clear();
-            // Close all chat streams
-            for (const [, stream] of this.chatStreams) {
-                stream.emit("event", { type: "error", error: "Connection closed" });
-                stream.close();
-            }
-            this.chatStreams.clear();
-            // Close all subscriptions
+        // Both "close" and "error" mean this socket is dead. ws typically emits
+        // "error" followed by "close", so route both through one handler and let the
+        // identity check dedupe: the first to fire nulls this.ws, the second no-ops.
+        const onDown = () => {
+            if (this.ws !== socket)
+                return;
+            this.handleDisconnect();
+        };
+        socket.on("close", onDown);
+        socket.on("error", onDown);
+    }
+    /**
+     * Reject in-flight requests and tear down the dead socket. If the close was
+     * unexpected (not an explicit `close()`) and auto-reconnect is enabled,
+     * schedule a reconnect that re-sends the active subscriptions.
+     */
+    handleDisconnect() {
+        this.dispatcherRunning = false;
+        this.ws = null;
+        // Reject all in-flight pending requests so callers don't hang forever.
+        for (const [, pending] of this.pendingRequests) {
+            if (pending.timer)
+                clearTimeout(pending.timer);
+            pending.reject(new Error("WebSocket connection closed"));
+        }
+        this.pendingRequests.clear();
+        if (this.registerToolsAck) {
+            this.registerToolsAck.reject(new Error("WebSocket connection closed"));
+            this.registerToolsAck = null;
+        }
+        // Close all chat streams (they are one-shot; not replayed on reconnect).
+        for (const [, stream] of this.chatStreams) {
+            stream.emit("event", { type: "error", error: "Connection closed" });
+            stream.close();
+        }
+        this.chatStreams.clear();
+        const shouldReconnect = this.autoReconnect && !this.closed && this.subscriptionParams.size > 0;
+        if (shouldReconnect) {
+            this.scheduleReconnect();
+        }
+        else {
+            // No reconnect: tear down subscriptions too.
             for (const [, stream] of this.subscriptions) {
                 stream.close();
             }
             this.subscriptions.clear();
-            this.ws = null;
-        });
+            this.subscriptionParams.clear();
+        }
+    }
+    /**
+     * Reconnect with capped exponential backoff + jitter, then re-send the
+     * subscribe messages for every active subscription so the SAME EventStream
+     * keeps delivering mutations after a transient drop.
+     */
+    scheduleReconnect() {
+        if (this.reconnecting)
+            return;
+        this.reconnecting = true;
+        const attempt = async () => {
+            // Bail if the client was closed, or if every subscription was torn down
+            // (e.g. unsubscribed) while a reconnect was in-flight — reconnect was only
+            // opted into because subscriptions existed, so there's nothing to restore.
+            if (this.closed || this.subscriptionParams.size === 0) {
+                this.reconnecting = false;
+                return;
+            }
+            if (this.reconnectMaxAttempts > 0 &&
+                this.reconnectAttempts >= this.reconnectMaxAttempts) {
+                // Give up: tear down subscriptions and notify consumers.
+                this.reconnecting = false;
+                for (const [, stream] of this.subscriptions) {
+                    stream.emit("error", "WebSocket reconnect failed");
+                    stream.close();
+                }
+                this.subscriptions.clear();
+                this.subscriptionParams.clear();
+                return;
+            }
+            const delay = this.computeBackoff(this.reconnectAttempts);
+            this.reconnectAttempts++;
+            await new Promise((r) => setTimeout(r, delay));
+            // Re-check after the backoff delay: close() or a full unsubscribe may have
+            // happened while we were waiting, in which case skip reopening the socket.
+            if (this.closed || this.subscriptionParams.size === 0) {
+                this.reconnecting = false;
+                return;
+            }
+            try {
+                // Route through ensureConnected() so a request-driven connect and this
+                // reconnect share one in-flight connectPromise/socket — opening two live
+                // sockets would misroute responses.
+                await this.ensureConnected();
+                // close() may have been called while the connect was in-flight; if so,
+                // tear down the freshly-opened socket instead of leaving it orphaned.
+                if (this.closed) {
+                    try {
+                        this.ws?.close?.();
+                    }
+                    catch {
+                        /* already closing */
+                    }
+                    this.ws = null;
+                    this.dispatcherRunning = false;
+                    this.reconnecting = false;
+                    return;
+                }
+                // Success — reset backoff and replay subscriptions.
+                this.reconnectAttempts = 0;
+                this.reconnecting = false;
+                await this.resubscribeAll();
+            }
+            catch {
+                // Connect failed — schedule the next attempt WITHOUT recursive await so
+                // a prolonged outage can't build an unbounded promise chain.
+                setTimeout(() => void attempt(), 0);
+            }
+        };
+        void attempt();
+    }
+    /** Re-send Subscribe frames for every tracked subscription after a reconnect. */
+    async resubscribeAll() {
+        for (const [collection, options] of this.subscriptionParams) {
+            const stream = this.subscriptions.get(collection);
+            if (!stream || stream.closed)
+                continue;
+            const messageId = this.genMessageId();
+            const request = {
+                type: "Subscribe",
+                messageId,
+                payload: {
+                    collection,
+                    ...(options?.filterField && { filter_field: options.filterField }),
+                    ...(options?.filterValue && { filter_value: options.filterValue }),
+                },
+            };
+            try {
+                await this.sendRequest(request);
+            }
+            catch {
+                // If the re-subscribe ack fails, leave it tracked; the next
+                // disconnect/reconnect cycle will attempt it again.
+            }
+        }
     }
     routeMessage(msg) {
         switch (msg.type) {
@@ -2085,15 +2356,7 @@ class WebSocketClient {
                     msg.payload?.messageId;
                 let matched = false;
                 if (messageId && this.pendingRequests.has(messageId)) {
-                    const pending = this.pendingRequests.get(messageId);
-                    this.pendingRequests.delete(messageId);
-                    if (msg.type === "Error") {
-                        pending.reject(new Error(msg.message || "Unknown error"));
-                    }
-                    else {
-                        pending.resolve(msg.payload);
-                    }
-                    matched = true;
+                    matched = this.settlePending(messageId, msg.type === "Error", msg);
                 }
                 if (!matched && this.registerToolsAck) {
                     const ack = this.registerToolsAck;
@@ -2106,19 +2369,14 @@ class WebSocketClient {
                     }
                     matched = true;
                 }
-                // Server doesn't echo messageId — if there's exactly one pending
+                // Server doesn't echo messageId at all — if there's exactly one pending
                 // request, deliver the response to it (sequential request/response).
-                if (!matched && this.pendingRequests.size === 1) {
-                    const entry = this.pendingRequests.entries().next().value;
-                    const key = entry[0];
-                    const pending = entry[1];
-                    this.pendingRequests.delete(key);
-                    if (msg.type === "Error") {
-                        pending.reject(new Error(msg.message || "Unknown error"));
-                    }
-                    else {
-                        pending.resolve(msg.payload);
-                    }
+                // Only when messageId is absent: a present-but-unmatched id means a late
+                // response for an already-settled/timed-out request, which must NOT be
+                // misrouted to whatever request happens to still be pending.
+                if (!matched && !messageId && this.pendingRequests.size === 1) {
+                    const key = this.pendingRequests.keys().next().value;
+                    this.settlePending(key, msg.type === "Error", msg);
                 }
                 break;
             }
@@ -2205,16 +2463,46 @@ class WebSocketClient {
         await this.ensureConnected();
         const messageId = request.messageId || request.message_id;
         return new Promise((resolve, reject) => {
-            this.pendingRequests.set(messageId, { resolve, reject });
+            // Per-request timeout: reject if no response arrives in the window so a
+            // dropped/never-answered response can't leave the promise pending forever.
+            let timer;
+            if (this.requestTimeoutMs > 0) {
+                timer = setTimeout(() => {
+                    if (this.pendingRequests.delete(messageId)) {
+                        reject(new Error(`WebSocket request "${request.type}" timed out after ${this.requestTimeoutMs}ms`));
+                    }
+                }, this.requestTimeoutMs);
+                // Don't keep the process alive just for this timer.
+                timer?.unref?.();
+            }
+            this.pendingRequests.set(messageId, { resolve, reject, timer });
             try {
                 this.ws.send(JSON.stringify(request));
             }
             catch (err) {
                 this.pendingRequests.delete(messageId);
+                if (timer)
+                    clearTimeout(timer);
                 reject(err);
             }
         });
     }
+    /** Resolve/reject a pending request, clearing its timeout timer. */
+    settlePending(messageId, isError, msg) {
+        const pending = this.pendingRequests.get(messageId);
+        if (!pending)
+            return false;
+        this.pendingRequests.delete(messageId);
+        if (pending.timer)
+            clearTimeout(pending.timer);
+        if (isError) {
+            pending.reject(new Error(msg.message || "Unknown error"));
+        }
+        else {
+            pending.resolve(msg.payload);
+        }
+        return true;
+    }
     /**
      * Find all records in a collection via WebSocket.
      */
@@ -2239,6 +2527,8 @@ class WebSocketClient {
         const messageId = this.genMessageId();
         const stream = new EventStream();
         this.subscriptions.set(collection, stream);
+        // Track params so the subscription can be replayed on reconnect.
+        this.subscriptionParams.set(collection, options);
         const request = {
             type: "Subscribe",
             messageId,
@@ -2254,10 +2544,23 @@ class WebSocketClient {
         }
         catch (err) {
             this.subscriptions.delete(collection);
+            this.subscriptionParams.delete(collection);
             throw err;
         }
         return stream;
     }
+    /**
+     * Unsubscribe from a collection's mutation notifications. This is an
+     * intentional teardown, so the subscription is NOT replayed on reconnect.
+     */
+    unsubscribe(collection) {
+        const stream = this.subscriptions.get(collection);
+        this.subscriptions.delete(collection);
+        this.subscriptionParams.delete(collection);
+        if (stream && !stream.closed) {
+            stream.close();
+        }
+    }
     /**
      * Send a chat message and receive a streaming response.
      * Returns an EventStream that emits "event" with ChatStreamEvent objects.
@@ -2481,8 +2784,42 @@ class WebSocketClient {
     }
     /**
      * Close the WebSocket connection.
+     *
+     * This is an INTENTIONAL close: it disables auto-reconnect, rejects any
+     * in-flight requests, and tears down all subscriptions/chat streams so
+     * nothing is replayed afterward.
      */
     close() {
+        // Mark intentional so the close handler doesn't trigger a reconnect.
+        this.closed = true;
+        this.reconnecting = false;
+        // Reject any in-flight requests and clear their timers.
+        for (const [, pending] of this.pendingRequests) {
+            if (pending.timer)
+                clearTimeout(pending.timer);
+            pending.reject(new Error("WebSocket connection closed"));
+        }
+        this.pendingRequests.clear();
+        // Tear down subscriptions + their replay bookkeeping.
+        for (const [, stream] of this.subscriptions) {
+            if (!stream.closed)
+                stream.close();
+        }
+        this.subscriptions.clear();
+        this.subscriptionParams.clear();
+        // Reject any in-flight tool registration ack. Done here (not just in the
+        // ws "close" handler) so it's cleaned up even when this.ws is already null.
+        if (this.registerToolsAck) {
+            this.registerToolsAck.reject(new Error("WebSocket connection closed"));
+            this.registerToolsAck = null;
+        }
+        // Tear down chat streams immediately; they are one-shot and not replayed,
+        // and we can't rely on the underlying ws "close" event having fired.
+        for (const [, stream] of this.chatStreams) {
+            stream.emit("event", { type: "error", error: "Connection closed" });
+            stream.close();
+        }
+        this.chatStreams.clear();
         if (this.ws) {
             this.ws.close();
             this.ws = null;